Update app.py

app.py

@@ -10,48 +10,55 @@ import pandas as pd
 import random
 import time
 
-# --- Main Configuration ---
+# --- Configuration Principale ---
 KRYPTO_LORA = {
     "repo": "Econogoat/Krypt0_LORA", 
     "trigger": "Krypt0",
     "adapter_name": "krypt0"
 }
 
-# Load prompts for the randomize button
-df = pd.read_csv('prompts.csv', header=None)
-prompt_values = df.values.flatten()
-
-# Get access token from Space secrets
-HF_TOKEN = os.getenv("HF_TOKEN")
-if not HF_TOKEN:
-    print("WARNING: HF_TOKEN secret is not set. Gated model downloads may fail.")
-
-# --- Model Initialization ---
-# 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.")
+# --- Lazy Loading Setup ---
+# CORRECTION : On ne charge rien ici. Les modèles sont initialisés à None.
+pipe = None
+good_vae = None
+taef1 = None
 dtype = torch.bfloat16
 base_model = "black-forest-labs/FLUX.1-dev"
 
-# 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.")
+# Charger les prompts pour le bouton de randomisation
+df = pd.read_csv('prompts.csv', header=None)
+prompt_values = df.values.flatten()
 
-# Load the LoRA adapter once on startup
-print(f"Loading on-board LoRA: {KRYPTO_LORA['repo']}")
-pipe.load_lora_weights(
-    KRYPTO_LORA['repo'],
-    low_cpu_mem_usage=True,
-    adapter_name=KRYPTO_LORA['adapter_name'],
-    token=HF_TOKEN
-)
-print("LoRA loaded successfully.")
+# Récupérer le token d'accès depuis les secrets du Space
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-MAX_SEED = 2**32 - 1
 
-# Monkey-patch the pipeline for live preview
-pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
+def load_models():
+    """
+    Cette fonction charge tous les modèles et les place sur le GPU.
+    Elle n'est appelée qu'une seule fois, lors de la première génération.
+    """
+    global pipe, good_vae, taef1
+    
+    print("Performing first-time setup: loading models to GPU...")
+    
+    # On charge tout directement sur CUDA car cette fonction est appelée depuis un contexte GPU.
+    taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to("cuda")
+    good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype, token=HF_TOKEN).to("cuda")
+    pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1, token=HF_TOKEN).to("cuda")
+    
+    print(f"Loading on-board LoRA: {KRYPTO_LORA['repo']}")
+    pipe.load_lora_weights(
+        KRYPTO_LORA['repo'],
+        low_cpu_mem_usage=True,
+        adapter_name=KRYPTO_LORA['adapter_name'],
+        token=HF_TOKEN
+    )
+    
+    # Monkey-patch the pipeline for live preview
+    pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
+    
+    print("Models loaded and ready on GPU.")
 
 
 def calculate_dimensions(aspect_ratio, resolution):
@@ -70,7 +77,6 @@ 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):
     generator = torch.Generator(device="cuda").manual_seed(seed)
     
@@ -102,37 +108,28 @@ def update_history(new_image, history):
     
 @spaces.GPU(duration=75)
 def run_generation(prompt, lora_scale, cfg_scale, steps, randomize_seed, seed, aspect_ratio, base_resolution, progress=gr.Progress(track_tqdm=True)):
+    global pipe
     if not prompt:
         raise gr.Error("Prompt cannot be empty.")
     
-    # 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.")
+    # CORRECTION : On charge les modèles seulement si ce n'est pas déjà fait.
+    if pipe is None:
+        load_models()
 
-        for image, progress_update in generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress):
-            yield image, seed, progress_update
+    prompt_mash = f"{KRYPTO_LORA['trigger']}, {prompt}"
+    print("Final prompt:", prompt_mash)
     
-    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()
+    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.")
 
+    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