import gradio as gr
from diffusers import DiffusionPipeline
import torch

# Load the Riffusion model on CPU
pipeline = DiffusionPipeline.from_pretrained("qualcomm/Riffusion")
pipeline = pipeline.to("cpu")

def generate_image(prompt, seed):
    # Set the random seed for reproducibility
    generator = torch.Generator().manual_seed(seed)
    
    # Generate the image
    image = pipeline(prompt=prompt, generator=generator).images[0]
    
    return image

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Enter a description of the image you want to generate"),
        gr.Number(label="Seed", value=42, precision=0)
    ],
    outputs=gr.Image(type="pil", label="Generated Image"),
    title="Riffusion Image Generator (CPU Version)",
    description="Generate images using the Riffusion model from Qualcomm (running on CPU)"
)

# Launch the app
iface.launch()