import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

model_id = "runwayml/stable-diffusion-v1-5"
device = "cuda" if torch.cuda.is_available() else "cpu"

# load once on startup
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to(device)
pipe.enable_attention_slicing()

def generate(prompt, steps, guidance, seed):
    generator = None
    if seed not in (None, "", "none"):
        generator = torch.Generator(device).manual_seed(int(seed))
    image = pipe(prompt, num_inference_steps=int(steps), guidance_scale=float(guidance), generator=generator).images[0]
    return image

with gr.Blocks() as demo:
    gr.Markdown("# Stable Diffusion — Space")
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", lines=2, value="A cinematic portrait of a Muslim scholar reading under a lamp, warm tones, detailed, realistic")
    with gr.Row():
        steps = gr.Slider(10, 50, value=30, step=1, label="Steps")
        guidance = gr.Slider(1.0, 12.0, value=7.5, step=0.5, label="Guidance")
        seed = gr.Textbox(label="Seed (optional)")
    btn = gr.Button("Generate")
    output = gr.Image(label="Generated image")

    btn.click(generate, inputs=[prompt, steps, guidance, seed], outputs=[output])

if __name__ == "__main__":
    demo.launch()