app.py

import spaces  # DOIT ETRE EN PREMIER !
import gradio as gr
import torch
from diffusers import FluxPipeline
import random

# Variable globale pour le pipeline
pipe = None

def load_model():
    """Charge le modele une seule fois"""
    global pipe
    if pipe is None:
        print("Chargement du modele...")
        pipe = FluxPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-dev",
            torch_dtype=torch.bfloat16
        )
        # Charger le LoRA
        pipe.load_lora_weights(
            "Heartsync/Flux-NSFW-uncensored",
            weight_name="lora.safetensors"
        )
        print("Modele pret!")
    return pipe

@spaces.GPU(duration=120)
def generate_image(prompt, seed=-1, steps=20, guidance=3.5, width=768, height=768):
    """Genere une image a partir d'un prompt"""

    # Charger le modele
    pipe = load_model()
    pipe.to("cuda")

    if seed == -1:
        seed = random.randint(0, 2147483647)

    generator = torch.Generator("cuda").manual_seed(seed)

    with torch.inference_mode():
        image = pipe(
            prompt=prompt,
            guidance_scale=guidance,
            num_inference_steps=int(steps),
            height=int(height),
            width=int(width),
            generator=generator
        ).images[0]

    return image, f"Seed: {seed}"


# Interface Gradio
demo = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Decris ton image...", lines=3),
        gr.Number(label="Seed (-1 = random)", value=-1),
        gr.Slider(10, 30, value=20, step=1, label="Steps"),
        gr.Slider(1, 10, value=3.5, step=0.5, label="Guidance Scale"),
        gr.Slider(512, 1024, value=768, step=64, label="Width"),
        gr.Slider(512, 1024, value=768, step=64, label="Height"),
    ],
    outputs=[
        gr.Image(label="Image generee"),
        gr.Textbox(label="Info")
    ],
    title="FLUX.1-dev Image Generator",
    description="Genere des images avec FLUX.1-dev + LoRA",
    api_name="generate"
)

demo.launch()