Update app.py

app.py

@@ -9,91 +9,71 @@ from diffusers.utils import load_image
 import pandas as pd
 import random
 import time
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
 # --- Configuration Principale ---
 KRYPTO_LORA = {"repo": "Econogoat/Krypt0_LORA", "trigger": "Krypt0", "adapter_name": "krypt0"}
-# --- CORRECTION DÉFINITIVE : Utilisation du bon ID de modèle, comme vous l'avez demandé ---
-GEMMA_MODEL_ID = "google/gemma-1.1-2b-it" 
 BASE_IMAGE_MODEL = "black-forest-labs/FLUX.1-dev"
 
-# --- Pré-chargement sur CPU des éléments légers UNIQUEMENT ---
-print("Pré-chargement des tokenizers sur CPU...")
-HF_TOKEN = os.getenv("HF_TOKEN")
-gemma_tokenizer = AutoTokenizer.from_pretrained(GEMMA_MODEL_ID, token=HF_TOKEN)
-print("Pré-chargement terminé. Les modèles lourds seront chargés à la demande sur le GPU.")
+# Charger les prompts
+df = pd.read_csv('prompts.csv', header=None)
+prompt_values = df.values.flatten()
 
-# --- Variables globales pour conserver les modèles en mémoire sur le GPU ---
+# --- Variables globales pour conserver les modèles en mémoire ---
+# On utilise une stratégie de chargement à la demande ("lazy loading") pour être compatible avec ZeroGPU
 pipe = None
 good_vae = None
-gemma_model = None
 
 MAX_SEED = 2**32 - 1
 
-@spaces.GPU(duration=180) # Durée augmentée pour le premier chargement
-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 pipe, good_vae, gemma_model
+@spaces.GPU(duration=180) # Durée pour accommoder le premier chargement du modèle d'image
+def run_generation(prompt, lora_scale, cfg_scale, steps, randomize_seed, seed, aspect_ratio, base_resolution, progress=gr.Progress(track_tqdm=True)):
+    global pipe, good_vae
     if not prompt:
         raise gr.Error("Prompt cannot be empty.")
 
     device_gpu = "cuda"
     device_cpu = "cpu"
     dtype = torch.bfloat16
-    final_prompt = prompt
     
-    # --- Chargement à la demande ("Lazy Loading") des modèles sur le GPU ---
+    # --- Chargement à la demande ("Lazy Loading") du modèle d'image sur le GPU ---
+    # Cette section ne s'exécute qu'une seule fois, lors du tout premier appel.
     if pipe is None:
         print("Premier appel : Chargement du pipeline d'image sur GPU...")
+        HF_TOKEN = os.getenv("HF_TOKEN")
         taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device_gpu)
         good_vae = AutoencoderKL.from_pretrained(BASE_IMAGE_MODEL, subfolder="vae", torch_dtype=dtype, token=HF_TOKEN).to(device_gpu)
         pipe = DiffusionPipeline.from_pretrained(BASE_IMAGE_MODEL, torch_dtype=dtype, vae=taef1, token=HF_TOKEN).to(device_gpu)
         print("Chargement du LoRA sur le pipeline GPU...")
         pipe.load_lora_weights(KRYPTO_LORA['repo'], low_cpu_mem_usage=False, adapter_name=KRYPTO_LORA['adapter_name'], token=HF_TOKEN)
         print("Pipeline d'image prêt.")
+        # Ajout de la méthode de prévisualisation
         pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
-    if enhance_prompt:
-        if gemma_model is None:
-            print(f"Premier appel d'amélioration : Chargement de {GEMMA_MODEL_ID} sur GPU...")
-            # La quantization est moins critique pour ce petit modèle, mais reste une bonne pratique.
-            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"
-            )
-            print("Modèle Gemma prêt.")
-        
-        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. Focus on describing the scene, the subject, the environment, the lighting, the colors, and a potential artistic style. Do not add any conversational text or refuse the request. Only output the enhanced prompt."
-        chat = [{"role": "user", "content": f"{system_prompt}\n\nUser idea: \"{user_prompt}\""}]
-        prompt_for_gemma = gemma_tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
-        inputs = gemma_tokenizer(prompt_for_gemma, return_tensors="pt").to(device_gpu)
-        
-        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}")
-
-    prompt_mash = f"{KRYPTO_LORA['trigger']}, {final_prompt}"
-    print("Prompt final envoyé au modèle d'image:", prompt_mash)
+    # Création du prompt final
+    prompt_mash = f"{KRYPTO_LORA['trigger']}, {prompt}"
+    print("Prompt final envoyé au modèle:", prompt_mash)
     
+    # Activation du LoRA
     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)
+    
+    # Appel du générateur d'image
     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,
+        prompt=prompt_mash,
+        num_inference_steps=steps,
+        guidance_scale=cfg_scale,
+        width=width,
+        height=height,
+        generator=generator,
+        output_type="pil",
+        good_vae=good_vae,
     )
 
     final_image = None
@@ -102,6 +82,14 @@ def run_generation(prompt, enhance_prompt, lora_scale, cfg_scale, steps, randomi
         progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {i + 1}; --total: {steps};"></div></div>'
         yield image, seed, gr.update(value=progress_bar, visible=True)
     
+    # Une fois la génération terminée, on peut libérer de la VRAM en déplaçant le modèle sur le CPU
+    # C'est optionnel mais une bonne pratique dans les environnements managés
+    print("Génération terminée. Déplacement du pipeline vers le CPU pour libérer la VRAM.")
+    pipe.to(device_cpu)
+    good_vae.to(device_cpu)
+    torch.cuda.empty_cache()
+    pipe = None # Force le rechargement au prochain appel
+    
     yield final_image, seed, gr.update(visible=False)
 
 def calculate_dimensions(aspect_ratio, resolution):
@@ -154,10 +142,6 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as app:
                         random_prompt_btn = gr.Button("Random Prompt", elem_id="random_prompt_btn")
                     with gr.Column(scale=5):
                         prompt = gr.Textbox(label="Prompt", lines=2, placeholder="e.g., a portrait of a warrior queen")
-                enhance_prompt_checkbox = gr.Checkbox(
-                    label="Improve prompt with AI", value=True, 
-                    info="Uses Gemma to automatically enrich your prompt with more details before generation."
-                )
             with gr.Group():
                 aspect_ratio = gr.Radio(
                     label="Aspect Ratio",
@@ -181,15 +165,20 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as app:
             result = gr.Image(label="Generated Image", interactive=False, show_share_button=True)
             with gr.Accordion("History", open=False):
                 history_gallery = gr.Gallery(label="History", columns=4, object_fit="contain", interactive=False)
+    
     def get_random_prompt():
         return random.choice(prompt_values)
+    
     random_prompt_btn.click(fn=get_random_prompt, inputs=[], outputs=[prompt])
+    
+    # Les entrées de la fonction de génération sont simplifiées
     generation_event = gr.on(
         triggers=[generate_button.click, prompt.submit],
         fn=run_generation,
-        inputs=[prompt, enhance_prompt_checkbox, lora_scale, cfg_scale, steps, randomize_seed, seed, aspect_ratio, base_resolution],
+        inputs=[prompt, lora_scale, cfg_scale, steps, randomize_seed, seed, aspect_ratio, base_resolution],
         outputs=[result, seed, progress_bar]
     )
+    
     generation_event.then(fn=update_history, inputs=[result, history_gallery], outputs=history_gallery)
 
 app.queue(max_size=20)