import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Load model without specifying fp16 revision
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16  # Still use fp16 precision
).to("cuda")

def generate_image(prompt, negative_prompt="", steps=30, guidance_scale=7.5):
    image = pipe(
        prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=steps,
        guidance_scale=guidance_scale
    ).images[0]
    return image

with gr.Blocks(title="RimageGen") as demo:
    gr.Markdown("## 🎨 Text-to-Image Generator")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", placeholder="A cute corgi wearing a crown")
            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="blurry, deformed")
            steps = gr.Slider(1, 50, value=30, label="Steps")
            guidance = gr.Slider(1, 15, value=7.5, label="Guidance Scale")
            submit = gr.Button("Generate")
        with gr.Column():
            output = gr.Image(label="Result")

    submit.click(generate_image, inputs=[prompt, negative_prompt, steps, guidance], outputs=output)

demo.launch()