import gradio as gr
from diffusers import DiffusionPipeline
import torch
import random

# Use lightweight model (faster & less resources)
model_id = "OFA-Sys/small-stable-diffusion-v0"  # 35x smaller than SDXL

# Load model with optimizations
pipe = DiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16
).to("cuda")

def generate(prompt):
    # Random seed for unique generations
    random_seed = random.randint(0, 2147483647)
    generator = torch.Generator("cuda").manual_seed(random_seed)
    
    # Generate image with variations
    image = pipe(
        prompt,
        num_inference_steps=20,  # Faster generation
        generator=generator
    ).images[0]
    
    return image

# Simple interface
gr.Interface(
    fn=generate,
    inputs=gr.Textbox(label="Enter text prompt"),
    outputs=gr.Image(label="Generated Image"),
    title="Simple AI Image Generator",
    description="Type anything - get random images every time!",
    allow_flagging="never"  # Remove feedback buttons
).launch()