Update app.py

app.py

@@ -12,7 +12,6 @@ import time
 
 # --- Main Configuration ---
 KRYPTO_LORA = {
-    # CORRECTION : Le nom du dépôt était mal orthographié (O majuscule au lieu d'un zéro).
     "repo": "Econogoat/Krypt0_LORA", 
     "trigger": "Krypt0",
     "adapter_name": "krypt0"
@@ -28,17 +27,16 @@ if not HF_TOKEN:
     print("WARNING: HF_TOKEN secret is not set. Gated model downloads may fail.")
 
 # --- Model Initialization ---
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Using device: {device}")
+# CORRECTION : On ne détecte plus le device ici, on charge tout sur CPU par défaut.
+print("Loading all models to CPU by default for ZeroGPU compatibility.")
 dtype = torch.bfloat16
 base_model = "black-forest-labs/FLUX.1-dev"
 
-# Load model components
-print("Loading model components...")
-taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
-good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype, token=HF_TOKEN).to(device)
-pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1, token=HF_TOKEN).to(device)
-print("Models loaded.")
+# CORRECTION : Tous les `.to(device)` sont retirés. Les modèles restent sur le CPU au démarrage.
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
+good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype, token=HF_TOKEN)
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1, token=HF_TOKEN)
+print("Models loaded on CPU.")
 
 # Load the LoRA adapter once on startup
 print(f"Loading on-board LoRA: {KRYPTO_LORA['repo']}")
@@ -46,7 +44,7 @@ pipe.load_lora_weights(
     KRYPTO_LORA['repo'],
     low_cpu_mem_usage=True,
     adapter_name=KRYPTO_LORA['adapter_name'],
-    token=HF_TOKEN # Ajout du token ici aussi pour les LoRA privés/protégés
+    token=HF_TOKEN
 )
 print("LoRA loaded successfully.")
 
@@ -72,9 +70,9 @@ def calculate_dimensions(aspect_ratio, resolution):
     height = (height // 64) * 64
     return width, height
 
+# CORRECTION : Cette fonction suppose maintenant que `pipe` est déjà sur le bon device.
 def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress):
-    pipe.to(device) 
-    generator = torch.Generator(device=device).manual_seed(seed)
+    generator = torch.Generator(device="cuda").manual_seed(seed)
     
     image_generator = pipe.flux_pipe_call_that_returns_an_iterable_of_images(
         prompt=prompt_mash,
@@ -106,21 +104,35 @@ def update_history(new_image, history):
 def run_generation(prompt, lora_scale, cfg_scale, steps, randomize_seed, seed, aspect_ratio, base_resolution, progress=gr.Progress(track_tqdm=True)):
     if not prompt:
         raise gr.Error("Prompt cannot be empty.")
-
-    prompt_mash = f"{KRYPTO_LORA['trigger']}, {prompt}"
-    print("Final prompt:", prompt_mash)
-    
-    pipe.set_adapters([KRYPTO_LORA['adapter_name']], adapter_weights=[lora_scale])
-    print(f"Adapter '{KRYPTO_LORA['adapter_name']}' activated with weight {lora_scale}.")
     
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+    # CORRECTION : C'est ici, à l'intérieur de la fonction @spaces.GPU, que l'on déplace les modèles sur le GPU.
+    print("Moving models to GPU for generation...")
+    pipe.to("cuda")
+    good_vae.to("cuda") # Il faut aussi déplacer le VAE de haute qualité
+
+    try:
+        prompt_mash = f"{KRYPTO_LORA['trigger']}, {prompt}"
+        print("Final prompt:", prompt_mash)
+        
+        pipe.set_adapters([KRYPTO_LORA['adapter_name']], adapter_weights=[lora_scale])
+        print(f"Adapter '{KRYPTO_LORA['adapter_name']}' activated with weight {lora_scale}.")
+        
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
+
+        width, height = calculate_dimensions(aspect_ratio, base_resolution)
+        print(f"Generating a {width}x{height} image.")
 
-    width, height = calculate_dimensions(aspect_ratio, base_resolution)
-    print(f"Generating a {width}x{height} image.")
+        for image, progress_update in generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress):
+            yield image, seed, progress_update
+    
+    finally:
+        # CORRECTION : Bonne pratique, on nettoie en déplaçant les modèles vers le CPU après usage.
+        print("Moving models back to CPU.")
+        pipe.to("cpu")
+        good_vae.to("cpu")
+        torch.cuda.empty_cache()
 
-    for image, progress_update in generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress):
-        yield image, seed, progress_update
 
 run_generation.zerogpu = True
 
@@ -180,7 +192,6 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as app:
                 )
 
             # Advanced Settings
-            # CORRECTION : L'accordéon doit être fermé par défaut.
             with gr.Accordion("Advanced Settings", open=False):
                 base_resolution = gr.Slider(label="Resolution (longest side)", minimum=768, maximum=1408, step=64, value=1024)
                 steps = gr.Slider(label="Generation Steps", minimum=4, maximum=50, step=1, value=20)