import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Load the model on CPU
model_id = "sergey-leshchenko/flux-realismo"
pipe = StableDiffusionPipeline.from_pretrained(model_id)
pipe.to("cpu")

def generate(prompt, guidance_scale=7.5, steps=25):
    image = pipe(prompt, guidance_scale=guidance_scale, num_inference_steps=steps).images[0]
    return image

# Gradio interface
gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Describe your image..."),
        gr.Slider(minimum=1, maximum=20, value=7.5, step=0.5, label="Guidance Scale"),
        gr.Slider(minimum=10, maximum=50, value=25, step=1, label="Inference Steps")
    ],
    outputs=gr.Image(type="pil"),
    title="Flux-Realismo Image Generator (CPU)",
    description="Generate realistic images using the Flux-Realismo model on CPU. May be slow."
).launch()