import gradio as gr
import torch
from diffusers import DiffusionPipeline

# Load the WAN 2.1 T2V Model
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = DiffusionPipeline.from_pretrained("Isi99999/Wan2.1-T2V-1.3B").to(device)

def generate_image(prompt):
    """Generates an image from text prompt using WAN 2.1"""
    image = pipe(prompt).images[0]
    return image

# Create Gradio UI
interface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Enter Prompt"),
    outputs=gr.Image(label="Generated Image"),
    title="WAN 2.1 - Text-to-Image Generation",
    description="Generate images from text using WAN 2.1 T2V model."
)

# Launch the app
interface.launch()