import torch
import gradio as gr
from diffusers import StableDiffusionPipeline

# Model path from Hugging Face
MODEL_NAME = "Evados/DiffSynth-Studio-Lora-Wan2.1-ComfyUI"
LORA_FILE = "Wan2.1-1.3b-lora-aesthetics-v1_new.safetensors"

def load_model():
    pipe = StableDiffusionPipeline.from_pretrained(MODEL_NAME, torch_dtype=torch.float16).to("cuda")
    pipe.load_lora_weights(f"{MODEL_NAME}/{LORA_FILE}")
    return pipe

pipe = load_model()

def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

# Gradio UI
iface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Enter your prompt"),
    outputs=gr.Image(label="Generated Image"),
    title="Wan2.1 LoRA Image Generator",
    description="Generate images using the Wan2.1 LoRA model. Enter a prompt to begin."
)

iface.launch()