Update app.py

app.py

@@ -22,47 +22,33 @@ prompt_values = df.values.flatten()
 # Récupérer le token
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# --- Initialisation des Modèles (sur CPU uniquement) ---
+# --- Initialisation des Modèles (sur CPU uniquement et SANS QUANTIZATION initiale) ---
 device_cpu = "cpu"
 dtype = torch.bfloat16
 base_model = "black-forest-labs/FLUX.1-dev"
 
-# --- CORRECTION DÉFINITIVE, BASÉE SUR VOTRE ANALYSE ---
-# Création d'une configuration BitsAndBytes qui spécifie explicitement `quant_type="nf4"`.
-# C'est la seule configuration supportée par bitsandbytes pour la quantization 4-bit sur CPU.
-bnb_config_cpu = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_compute_dtype=dtype
-)
-
-print(f"Chargement du LLM {GEMMA_MODEL_ID} sur CPU avec la config NF4...")
+# --- STRATÉGIE CORRIGÉE ---
+# On charge Gemma sur le CPU SANS le quantizer au démarrage pour éviter le conflit avec l'environnement de Spaces.
+# La quantization sera appliquée plus tard, uniquement sur le GPU.
+print(f"Chargement du tokenizer pour {GEMMA_MODEL_ID}...")
 gemma_tokenizer = AutoTokenizer.from_pretrained(GEMMA_MODEL_ID, token=HF_TOKEN)
-gemma_model = AutoModelForCausalLM.from_pretrained(
-    GEMMA_MODEL_ID,
-    quantization_config=bnb_config_cpu, # Utilisation de la configuration corrigée
-    token=HF_TOKEN,
-    device_map={'':device_cpu}
-)
-print("Modèle Gemma chargé.")
-
-# Le reste du chargement est correct
-print("Chargement des composants du modèle d'image sur CPU...")
+print("Chargement du modèle d'image sur CPU...")
 taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device_cpu)
 good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype, token=HF_TOKEN).to(device_cpu)
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1, token=HF_TOKEN).to(device_cpu)
-print("Modèles d'image chargés.")
-
-print(f"Chargement du LoRA : {KRYPTO_LORA['repo']}")
+print("Chargement du LoRA...")
 pipe.load_lora_weights(KRYPTO_LORA['repo'], low_cpu_mem_usage=True, adapter_name=KRYPTO_LORA['adapter_name'], token=HF_TOKEN)
-print("LoRA chargé.")
+print("Tous les modèles sont pré-chargés sur CPU.")
 
 MAX_SEED = 2**32 - 1
 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
-@spaces.GPU(duration=120)
+# On garde une référence globale pour ne recharger le modèle qu'une fois.
+gemma_model = None
+
+@spaces.GPU(duration=180) # Augmentation de la durée pour accommoder le chargement complet de Gemma
 def run_generation(prompt, enhance_prompt, lora_scale, cfg_scale, steps, randomize_seed, seed, aspect_ratio, base_resolution, progress=gr.Progress(track_tqdm=True)):
+    global gemma_model
     if not prompt:
         raise gr.Error("Prompt cannot be empty.")
 
@@ -70,10 +56,24 @@ def run_generation(prompt, enhance_prompt, lora_scale, cfg_scale, steps, randomi
     final_prompt = prompt
     
     if enhance_prompt:
-        print("Déplacement de Gemma sur le GPU...")
-        gemma_model.to(device_gpu)
-        print(f"Amélioration du prompt '{prompt}' avec Gemma...")
+        # --- CHARGEMENT DYNAMIQUE SUR GPU ---
+        if gemma_model is None:
+            print(f"Premier appel : Chargement de {GEMMA_MODEL_ID} sur GPU avec quantization 4-bit...")
+            bnb_config_gpu = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_compute_dtype=dtype
+            )
+            gemma_model = AutoModelForCausalLM.from_pretrained(
+                GEMMA_MODEL_ID,
+                quantization_config=bnb_config_gpu,
+                token=HF_TOKEN,
+                device_map="auto" # "auto" fonctionnera car on est DÉJÀ sur un environnement GPU
+            )
+            print("Modèle Gemma chargé sur GPU.")
         
+        print(f"Amélioration du prompt '{prompt}' avec Gemma...")
         system_prompt = (
             "You are an expert prompt engineer for a text-to-image AI. "
             "Your task is to take a user's simple idea and transform it into a rich, detailed, and visually descriptive prompt. "
@@ -87,11 +87,7 @@ def run_generation(prompt, enhance_prompt, lora_scale, cfg_scale, steps, randomi
         outputs = gemma_model.generate(**inputs, max_new_tokens=150, do_sample=True, temperature=0.7)
         input_length = inputs["input_ids"].shape[1]
         final_prompt = gemma_tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True).strip()
-
         print(f"Prompt amélioré : {final_prompt}")
-        print("Libération de la mémoire de Gemma (déplacement vers CPU)...")
-        gemma_model.to(device_cpu)
-        torch.cuda.empty_cache()
 
     prompt_mash = f"{KRYPTO_LORA['trigger']}, {final_prompt}"
     print("Prompt final envoyé au modèle d'image:", prompt_mash)
@@ -99,17 +95,13 @@ def run_generation(prompt, enhance_prompt, lora_scale, cfg_scale, steps, randomi
     print("Déplacement du pipeline d'image sur le GPU...")
     pipe.to(device_gpu)
     good_vae.to(device_gpu)
-
     pipe.set_adapters([KRYPTO_LORA['adapter_name']], adapter_weights=[lora_scale])
     
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
+    if randomize_seed: seed = random.randint(0, MAX_SEED)
     width, height = calculate_dimensions(aspect_ratio, base_resolution)
     print(f"Génération d'une image de {width}x{height} pixels.")
     
     generator = torch.Generator(device=device_gpu).manual_seed(seed)
-    
     image_generator = pipe.flux_pipe_call_that_returns_an_iterable_of_images(
         prompt=prompt_mash, num_inference_steps=steps, guidance_scale=cfg_scale,
         width=width, height=height, generator=generator, output_type="pil", good_vae=good_vae,