import torch
import gradio as gr
from diffusers import DiffusionPipeline

MODEL_ID = "black-forest-labs/FLUX.1-schnell"

# Pipeline laden (CPU)
pipe = DiffusionPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32
)
pipe.to("cpu")

def generate(prompt):
    image = pipe(
        prompt=prompt,
        num_inference_steps=20,
        guidance_scale=3.5
    ).images[0]
    return image

demo = gr.Interface(
    fn=generate,
    inputs=gr.Textbox(
        label="Prompt",
        placeholder="anime girl with fennec ears sitting on a log in the woods"
    ),
    outputs=gr.Image(type="pil"),
    title="FLUX Text-to-Image (CPU Space)",
    description="Runs on CPU using diffusers. Slow but works."
)

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