import gradio as gr
import torch
from diffusers import DiffusionPipeline

# Load the WAN 2.1 T2V Model

from diffusers import DiffusionPipeline

from diffusers import DiffusionPipeline

from diffusers import DiffusionPipeline

from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("sarthak247/Wan2.1-T2V-1.3B-nf4")

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(prompt).images[0]
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()