app.py

import spaces
import json
import os
import torch
import gradio as gr

from diffusers import AutoPipelineForText2Image
from PIL import PngImagePlugin

# -----------------------
# Config modèles
# -----------------------
BASE_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"

# LoRA "wrong" (améliore qualité + adhérence prompt)
LORA_WRONG_REPO = "minimaxir/sdxl-wrong-lora"  # [web:183]
LORA_WRONG_ADAPTER = "wrong_prompt"

# LoRA visages / détails
LORA_FACE_REPO = "akash-guptag/Detailers_By_Stable_Yogi"  
LORA_FACE_ADAPTER = "face_detail"

LORA_SEXY_REPO = "ntc-ai/SDXL-LoRA-slider.sexy"
LORA_SEXY_ADAPTER ="sexy"

LORA_DETAILLED_REPO = "ntc-ai/SDXL-LoRA-slider.extremely-detailed"
LORA_DETAILLED_ADAPTER = "extremely detailed"


DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32

# -----------------------
# Chargement pipeline
# -----------------------
print("Chargement SDXL...")
pipe = AutoPipelineForText2Image.from_pretrained(
    BASE_MODEL_ID,
    torch_dtype=DTYPE,
    variant="fp16" if DTYPE == torch.float16 else None,
    safety_checker=None,            # on désactive officiellement
    requires_safety_checker=False,
)
pipe.to(DEVICE)

# Monkey-patch du safety checker (NSFW fully bypass)
def dummy_safety_checker(images, **kwargs):
    # renvoie toujours "tout va bien"
    return images, [False] * len(images)

pipe.safety_checker = dummy_safety_checker
pipe.set_progress_bar_config(disable=True)

print("Chargement LoRA WRONG (prompt/qualité)...")
pipe.load_lora_weights(LORA_WRONG_REPO, adapter_name=LORA_WRONG_ADAPTER)

print("Chargement LoRA face/detail (Stable Yogi)...")
pipe.load_lora_weights(LORA_FACE_REPO, adapter_name=LORA_FACE_ADAPTER)

pipe.load_lora_weights(LORA_SEXY_REPO, adapter_name=LORA_SEXY_ADAPTER)

pipe.load_lora_weights(LORA_DETAILLED_REPO, adapter_name=LORA_DETAILLED_ADAPTER)


os.makedirs("outputs", exist_ok=True)


# -----------------------
# Fonction de génération
# -----------------------
@spaces.GPU()
def generate(
    prompt: str,
    negative: str,
    seed_in,
    steps: float,
    guidance: float,
    width: float,
    height: float,
    face_weight: float,
    script_name: str,
):
    # Seed robuste (Gradio envoie parfois int, parfois str)
    try:
        if isinstance(seed_in, str):
            seed_in = seed_in.strip()
            seed = int(float(seed_in)) if seed_in != "" else -1
        elif seed_in is None:
            seed = -1
        else:
            seed = int(seed_in)
    except Exception:
        seed = -1

    if seed >= 0:
        generator = torch.Generator(device=DEVICE).manual_seed(seed)
    else:
        generator = torch.Generator(device=DEVICE)

    # Prompt principal (on enrobe pour aider SDXL)
    final_prompt = (
        "masterpiece, best quality, extremely detailed, dynamic pose,cinematic "
        + (prompt or "1girl, caucasian, medium breasts, dressed like a princess,medium breasts, detailed face, realistic skin")
        + ", sharp focus"
    )

    # Negative : "wrong" est la clé pour sdxl-wrong-lora 
    if negative and negative.strip():
        final_negative = "wrong, " + negative.strip()
    else:
        final_negative = "wrong, blurry, low quality, deformed, bad anatomy, extra limbs"

    # On combine les deux LoRA : WRONG à 1.0, face en slider
        adapters = [LORA_WRONG_ADAPTER, LORA_FACE_ADAPTER, LORA_SEXY_ADAPTER, LORA_DETAILLED_ADAPTER]
    weights = [1.0, float(face_weight), 1.5, 2.0]
    pipe.set_adapters(adapters, adapter_weights=weights)

    try:
        result = pipe(
            prompt=final_prompt,
            negative_prompt=final_negative,
            num_inference_steps=int(steps),
            guidance_scale=float(guidance),
            width=int(width),
            height=int(height),
            generator=generator,
        )
    except Exception as e:
        return None, f"Erreur pendant la génération : {repr(e)}", ""

    image = result.images[0]

    metadata = {
        "base_model": BASE_MODEL_ID,
        "prompt_raw": prompt,
        "prompt_final": final_prompt,
        "negative_raw": negative,
        "negative_final": final_negative,
        "seed": seed,
        "steps": int(steps),
        "guidance": float(guidance),
        "width": int(width),
        "height": int(height),
        "lora_wrong_repo": LORA_WRONG_REPO,
        "lora_wrong_adapter": LORA_WRONG_ADAPTER,
        "lora_wrong_weight": 1.0,
        "lora_face_repo": LORA_FACE_REPO,
        "lora_face_adapter": LORA_FACE_ADAPTER,
        "lora_face_weight": float(face_weight),
    }

    base_name = script_name.strip().replace(" ", "_") if script_name else "sdxl_wrong_lora"
    img_path = os.path.join("outputs", f"{base_name}.png")
    json_path = os.path.join("outputs", f"{base_name}.json")

    pnginfo = PngImagePlugin.PngInfo()
    pnginfo.add_text("generation_params", json.dumps(metadata, ensure_ascii=False))
    image.save(img_path, pnginfo=pnginfo)

    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(metadata, f, ensure_ascii=False, indent=2)

    script_txt = json.dumps(metadata, ensure_ascii=False, indent=2)
    return image, script_txt, json_path


# -----------------------
# UI Gradio
# -----------------------
with gr.Blocks(title="SDXL + WRONG LoRA + Face Detail") as demo:
    gr.Markdown(
        "## SDXL 1.0 + LoRA **WRONG** (meilleure compréhension) + LoRA détail visage  \n"
        "- Utilise le LoRA `sdxl-wrong-lora` pour améliorer qualité et adhérence au prompt.[web:183][web:194]\n"
        "- Utilise `Detailers_By_Stable_Yogi` pour les visages/détails.[web:60][web:63]\n"
        "- NSFW débloqué (safety checker bypassé)."
    )

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="1girl nude, red dress, on the beach at sunset, cinematic lighting, smiling at viewer",
                value=None,
                lines=4,
            )
            negative = gr.Textbox(
                label="Negative prompt (\"wrong\" sera ajouté automatiquement)",
                placeholder="blurry, deformed, ugly, extra limbs, bad anatomy",
                value=None,
                lines=3,
            )

            seed = gr.Number(
                label="Seed (-1 = random)",
                value=-1,
            )
            steps = gr.Slider(
                minimum=20,
                maximum=60,
                value=40,
                step=1,
                label="Steps (plus haut = meilleure adhérence)",
            )
            guidance = gr.Slider(
                minimum=5.0,
                maximum=15.0,
                value=9.0,
                step=0.5,
                label="CFG / Guidance scale",
            )

            width = gr.Slider(
                minimum=512,
                maximum=1536,
                value=1024,
                step=64,
                label="Width (SDXL natif 1024)",
            )
            height = gr.Slider(
                minimum=512,
                maximum=1536,
                value=1024,
                step=64,
                label="Height (SDXL natif 1024)",
            )

            face_weight = gr.Slider(
                minimum=0.0,
                maximum=1.2,
                value=0.7,
                step=0.05,
                label="Force LoRA face/detail (Stable Yogi)",
            )

            script_name = gr.Textbox(
                label="Nom base pour l'image / script",
                value="sdxl_wrong_example",
            )

            run_btn = gr.Button("🚀 Générer", variant="primary")

        with gr.Column():
            out_img = gr.Image(
                label="Image générée (SDXL + WRONG + Face)",
            )
            out_script = gr.Textbox(
                label="Metadata / Script JSON",
                lines=20,
            )
            out_file = gr.File(
                label="Fichier JSON des paramètres (téléchargeable)",
            )

    run_btn.click(
        fn=generate,
        inputs=[
            prompt,
            negative,
            seed,
            steps,
            guidance,
            width,
            height,
            face_weight,
            script_name,
        ],
        outputs=[out_img, out_script, out_file],
    )

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