import gradio as gr
from diffusers import Flux2Pipeline
from peft import PeftModel
import torch, os

BASE_MODEL = "black-forest-labs/FLUX.2-klein-base-9B"

LORA_MODEL = "cyrildiagne/flux2-klein9b-lora-mlsharp-3d-repair"

# Load on CPU or balanced device_map
pipe = Flux2Pipeline.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float16,
    device_map="balanced",  # supported strategy (no crash)
)
pipe = PeftModel.from_pretrained(pipe, LORA_MODEL)

def run(prompt):
    # Move the model to GPU when this function runs
    pipe.to("cuda")

    image = pipe(
        prompt,
        num_inference_steps=25,
        guidance_scale=3.5,
    ).images[0]
    return image

demo = gr.Interface(fn=run, inputs="text", outputs="image")
demo.launch()