import gradio as gr
import torch
from diffusers import FluxPipeline
from PIL import Image
import os

# Chargement du modèle FLUX + LoRA Ghost Mannequin (le meilleur actuellement)
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-base-9B",   # version plus légère et rapide que FLUX.1-dev
    torch_dtype=torch.bfloat16
)
pipe.load_lora_weights("nhathoangfoto/FLUX.2-klein-ghost-mannequin")
pipe.to("cuda" if torch.cuda.is_available() else "cpu")

def generate_ghost_mannequin(image: Image.Image, strength: float = 0.75, steps: int = 28):
    if image is None:
        return None, "Veuillez uploader une image du t-shirt à plat."
    
    # Prompt optimisé pour ton cas exact
    prompt = (
        "3Dghostmannequin, t-shirt floating on invisible ghost mannequin, "
        "pure white background #FFFFFF, no shadows, no visible mannequin, "
        "clean professional product photography, realistic fabric texture, "
        "high detail, studio lighting, sharp focus"
    )
    
    negative_prompt = (
        "person, visible mannequin, head, neck, arms, hands, legs, shadow, "
        "background, wrinkles, dirt, text, logo too big, low quality, blurry"
    )
    
    # Génération
    result = pipe(
        prompt=prompt,
        image=image,
        strength=strength,           # 0.6 à 0.85 → plus la valeur est haute, plus on transforme l’image originale
        guidance_scale=3.5,
        num_inference_steps=steps,
        height=1024,
        width=1024,
        max_sequence_length=512
    ).images[0]
    
    return result, "✅ Image générée avec succès !"

# Interface Gradio
with gr.Blocks(title="Ghost Mannequin T-Shirt", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🧥 Ghost Mannequin T-Shirt\n**Transforme ton t-shirt à plat en photo pro sur fond blanc sans ombre**")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                type="pil",
                label="📸 Photo du t-shirt à plat (flat lay)",
                height=400
            )
            strength_slider = gr.Slider(
                minimum=0.5, maximum=0.95, value=0.75, step=0.05,
                label="Force de transformation (plus haut = plus de changement)"
            )
            steps_slider = gr.Slider(
                minimum=15, maximum=40, value=28, step=1,
                label="Nombre d'étapes (plus = meilleure qualité, plus lent)"
            )
            btn = gr.Button("🚀 Générer le Ghost Mannequin", variant="primary", size="large")
        
        with gr.Column():
            output_image = gr.Image(
                label="✅ Résultat : T-shirt sur fond blanc (mannequin invisible)",
                height=400
            )
            status = gr.Textbox(label="Statut", interactive=False)
    
    gr.Examples(
        examples=[["examples/tshirt_flat1.jpg"], ["examples/tshirt_flat2.jpg"]],
        inputs=input_image,
        label="Exemples de t-shirts à plat"
    )
    
    btn.click(
        fn=generate_ghost_mannequin,
        inputs=[input_image, strength_slider, steps_slider],
        outputs=[output_image, status]
    )

demo.launch()