Update app.py

app.py

@@ -16,77 +16,83 @@ from transformers import AutoModelForImageSegmentation
 from torchvision import transforms
 from pipeline import InstantCharacterFluxPipeline
 
-# =========================
-# CONFIG
-# =========================
+# =====================================================
+# CONFIG GÉNÉRALE
+# =====================================================
 MAX_SEED = np.iinfo(np.int32).max
-device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.float16  # L4: FP16 OK
-
+dtype = torch.float16  # parfait sur L4 (24 Go)
 HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACEHUB_API_TOKEN")
+
 def need_token_guard():
     if HF_TOKEN is None:
         raise gr.Error(
             "⚠️ Token manquant : ajoute un secret 'HF_TOKEN' (Settings → Repository secrets) "
-            "avec accès à black-forest-labs/FLUX.1-dev."
+            "avec ton token Hugging Face ayant accès à black-forest-labs/FLUX.1-dev."
         )
 
-# =========================
-# PATHS / WEIGHTS
-# =========================
+# =====================================================
+# TÉLÉCHARGEMENT DES PONTS ET MODÈLES
+# =====================================================
 base_model = "black-forest-labs/FLUX.1-dev"
 image_encoder_path = "google/siglip-so400m-patch14-384"
-# 🔻 On supprime l'encodeur 2 (DINOv2-giant) pour sauver ~10 Go RAM
-image_encoder_2_path = None
+image_encoder_2_path = "facebook/dinov2-giant"
 birefnet_path = "ZhengPeng7/BiRefNet"
 
-# Ton LoRA One Piece (local)
-onepiece_flux_lora_path = "./onepiece_flux_v2.safetensors"
-onepiece_flux_trigger = "onepiece style"
-
 def _dl(repo_id, filename, token=None):
     return hf_hub_download(repo_id=repo_id, filename=filename, token=token)
 
 need_token_guard()
-# Uniquement l'IP-Adapter (nécessaire à l'identité) — 5.6 Go
 ip_adapter_path = _dl("tencent/InstantCharacter", "instantcharacter_ip-adapter.bin", HF_TOKEN)
 
-# =========================
-# PIPELINE (GPU only, low RAM peak)
-# =========================
+makoto_style_lora_path = _dl("InstantX/FLUX.1-dev-LoRA-Makoto-Shinkai",
+                             "Makoto_Shinkai_style.safetensors", HF_TOKEN)
+ghibli_style_lora_path = _dl("InstantX/FLUX.1-dev-LoRA-Ghibli",
+                             "ghibli_style.safetensors", HF_TOKEN)
+
+# >>> Ton LoRA One Piece (FLUX) <<<
+onepiece_flux_lora_path = "./onepiece_flux_v2.safetensors"
+onepiece_flux_trigger = "onepiece style"
+
+# =====================================================
+# INITIALISATION DU PIPELINE (optimisée VRAM)
+# =====================================================
 pipe = InstantCharacterFluxPipeline.from_pretrained(
-    base_model,
-    torch_dtype=dtype,
-    token=HF_TOKEN,
-    low_cpu_mem_usage=True,   # ↓ pic RAM à l'init
+    base_model, torch_dtype=dtype, token=HF_TOKEN
 )
-pipe.to(device)
 
 try:
     pipe.enable_xformers_memory_efficient_attention()
 except Exception:
     pass
+try:
+    pipe.enable_model_cpu_offload()  # offload auto GPU/CPU
+except Exception:
+    try:
+        pipe.enable_sequential_cpu_offload()
+    except Exception:
+        pass
 
 pipe.set_progress_bar_config(disable=True)
 if hasattr(pipe, "vae"):
-    if hasattr(pipe.vae, "enable_slicing"): pipe.vae.enable_slicing()
-    if hasattr(pipe.vae, "enable_tiling"):  pipe.vae.enable_tiling()
+    if hasattr(pipe.vae, "enable_slicing"):
+        pipe.vae.enable_slicing()
+    if hasattr(pipe.vae, "enable_tiling"):
+        pipe.vae.enable_tiling()
 
-# 🔻 Init de l'adapter: 1 seul image encoder (SigLIP) + moins de tokens
-adapter_kwargs = dict(
+pipe.init_adapter(
     image_encoder_path=image_encoder_path,
-    subject_ipadapter_cfg=dict(subject_ip_adapter_path=ip_adapter_path, nb_token=512),  # 1024 -> 512 pour baisser mémoire
+    image_encoder_2_path=image_encoder_2_path,
+    subject_ipadapter_cfg=dict(subject_ip_adapter_path=ip_adapter_path, nb_token=1024),
 )
-# N'ajoute image_encoder_2_path que s'il existe
-if image_encoder_2_path:
-    adapter_kwargs["image_encoder_2_path"] = image_encoder_2_path
 
-pipe.init_adapter(**adapter_kwargs)
+# =====================================================
+# MATTEUR (BiRefNet) – SUR CPU POUR ÉCONOMISER LA VRAM
+# =====================================================
+birefnet = AutoModelForImageSegmentation.from_pretrained(
+    birefnet_path, trust_remote_code=True, token=HF_TOKEN
+)
+birefnet.to("cpu").eval()
 
-# =========================
-# BiRefNet : lazy-load sur CPU
-# =========================
-birefnet = None
 birefnet_transform_image = transforms.Compose([
     transforms.Resize((1024, 1024)),
     transforms.ToTensor(),
@@ -94,13 +100,6 @@ birefnet_transform_image = transforms.Compose([
 ])
 
 def remove_bkg(subject_image):
-    global birefnet
-    if birefnet is None:
-        birefnet = AutoModelForImageSegmentation.from_pretrained(
-            birefnet_path, trust_remote_code=True, token=HF_TOKEN
-        )
-        birefnet.to("cpu").eval()
-
     def infer_matting(img_pil):
         imgs = birefnet_transform_image(img_pil).unsqueeze(0).to("cpu")
         with torch.no_grad():
@@ -122,8 +121,10 @@ def remove_bkg(subject_image):
         if H == W: return image
         pad = abs(H - W)
         pad1, pad2 = pad // 2, pad - pad // 2
-        if H > W: pad_param = ((0,0),(pad1,pad2),(0,0))
-        else:     pad_param = ((pad1,pad2),(0,0),(0,0))
+        if H > W:
+            pad_param = ((0, 0), (pad1, pad2), (0, 0))
+        else:
+            pad_param = ((pad1, pad2), (0, 0), (0, 0))
         return np.pad(image, pad_param, "constant", constant_values=pad_value)
 
     mask = infer_matting(subject_image)[..., 0]
@@ -137,15 +138,15 @@ def remove_bkg(subject_image):
     crop = pad_to_square(crop)
     return Image.fromarray(crop.astype(np.uint8))
 
-# =========================
-# UTILS
-# =========================
+# =====================================================
+# OUTILS
+# =====================================================
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     return random.randint(0, MAX_SEED) if randomize_seed else seed
 
-# =========================
-# INFERENCE
-# =========================
+# =====================================================
+# GÉNÉRATION D'IMAGE
+# =====================================================
 @spaces.GPU
 def create_image(
     input_image,
@@ -154,54 +155,55 @@ def create_image(
     guidance_scale,
     num_inference_steps,
     seed,
-    use_onepiece_lora=True,
+    style_mode=None,
     lora_strength=0.85,
-    width=768,
-    height=768,
+    width=896,
+    height=896,
 ):
     if input_image is None:
         raise gr.Error("Merci d'uploader une image de visage.")
-    if use_onepiece_lora and not os.path.exists(onepiece_flux_lora_path):
+    if style_mode == "One Piece (FLUX LoRA)" and not os.path.exists(onepiece_flux_lora_path):
         raise gr.Error(f"Fichier LoRA manquant : {onepiece_flux_lora_path}")
 
     input_image = remove_bkg(input_image)
-    generator = torch.Generator(device=device).manual_seed(int(seed))
-
-    if use_onepiece_lora:
-        images = pipe.with_style_lora(
-            lora_file_path=onepiece_flux_lora_path,
-            trigger=onepiece_flux_trigger,
-            prompt=prompt,
-            num_inference_steps=int(num_inference_steps),
-            guidance_scale=float(guidance_scale),
-            width=int(width),
-            height=int(height),
-            subject_image=input_image,
-            subject_scale=float(scale),
-            lora_scale=float(lora_strength),
-            generator=generator,
-        ).images
+    generator = None  # évite conflits avec offload auto
+
+    if style_mode == "Makoto Shinkai style":
+        lora_file_path, trigger = makoto_style_lora_path, "Makoto Shinkai style"
+    elif style_mode == "Ghibli style":
+        lora_file_path, trigger = ghibli_style_lora_path, "ghibli style"
+    elif style_mode == "One Piece (FLUX LoRA)":
+        lora_file_path, trigger = onepiece_flux_lora_path, onepiece_flux_trigger
     else:
-        images = pipe(
-            prompt=prompt,
-            num_inference_steps=int(num_inference_steps),
-            guidance_scale=float(guidance_scale),
-            width=int(width),
-            height=int(height),
-            subject_image=input_image,
-            subject_scale=float(scale),
-            generator=generator,
-        ).images
-
+        lora_file_path, trigger = None, None
+
+    fn = pipe.with_style_lora if lora_file_path else pipe
+    kwargs = dict(
+        prompt=prompt,
+        num_inference_steps=int(num_inference_steps),
+        guidance_scale=float(guidance_scale),
+        width=int(width),
+        height=int(height),
+        subject_image=input_image,
+        subject_scale=float(scale),
+        generator=generator,
+    )
+    if lora_file_path:
+        kwargs.update(dict(
+            lora_file_path=lora_file_path,
+            trigger=trigger,
+            lora_scale=float(lora_strength),
+        ))
+    images = fn(**kwargs).images
     return images
 
-# =========================
-# UI
-# =========================
-title = "<h1 align='center'>InstantCharacter (FLUX.1-dev) + One Piece (FLUX LoRA) — single encoder</h1>"
+# =====================================================
+# INTERFACE GRADIO
+# =====================================================
+title = "<h1 align='center'>InstantCharacter (FLUX.1-dev) + One Piece (FLUX LoRA)</h1>"
 description = (
-    "GPU-only (FP16), low_cpu_mem_usage=True, **sans DINOv2-giant** pour éviter la limite RAM 30 Go. "
-    "Départ en 768×768, tu peux monter à 896→1024 si stable."
+    "<b>GPU :</b> Nvidia L4 24 Go — configuration optimisée VRAM.<br>"
+    "Résolution par défaut : 896 × 896 (monte à 1024 si stable)."
 )
 
 block = gr.Blocks(css="footer {visibility: hidden}").queue(concurrency_count=1, max_size=5, api_open=False)
@@ -216,17 +218,17 @@ with block:
                 value="onepiece style, a pirate character standing on a ship deck, shonen manga, strong black line art, cel shading, expressive eyes, dynamic pose, clean linework"
             )
             scale = gr.Slider(0.0, 1.5, 1.0, 0.01, label="Scale (face strength)")
-            use_onepiece_lora = gr.Checkbox(value=True, label="Use One Piece (FLUX LoRA)")
-            lora_strength = gr.Slider(0.0, 1.5, 0.85, 0.05, label="LoRA strength")
-
+            style_mode = gr.Dropdown(
+                ["None", "Makoto Shinkai style", "Ghibli style", "One Piece (FLUX LoRA)"],
+                value="One Piece (FLUX LoRA)",
+                label="Style",
+            )
+            lora_strength = gr.Slider(0.0, 1.5, 0.85, 0.05, label="LoRA strength (One Piece)")
             with gr.Accordion("Advanced Options", open=False):
                 guidance_scale = gr.Slider(1.0, 7.0, 3.5, 0.1, label="Guidance scale")
                 num_inference_steps = gr.Slider(5, 50, 28, 1, label="Inference steps")
                 seed = gr.Slider(-MAX_SEED, MAX_SEED, 123456, 1, label="Seed")
                 randomize_seed = gr.Checkbox(value=True, label="Randomize seed")
-                width  = gr.Slider(640, 1152, 768, 32, label="Width")
-                height = gr.Slider(640, 1152, 768, 32, label="Height")
-
             generate_button = gr.Button("Generate Image", variant="primary")
 
         with gr.Column():
@@ -240,7 +242,7 @@ with block:
     ).then(
         fn=create_image,
         inputs=[image_pil, prompt, scale, guidance_scale, num_inference_steps,
-                seed, use_onepiece_lora, lora_strength, width, height],
+                seed, style_mode, lora_strength],
         outputs=output_gallery,
     )