Update app.py

app.py

@@ -4,167 +4,208 @@ import torch
 from PIL import Image
 import spaces
 from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL
-from transformers import AutoModelForImageTextToText, AutoProcessor
 from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
 from diffusers.utils import load_image
 import pandas as pd
 import random
 import time
 
-# --- Configuration (statique) ---
-KRYPTO_LORA = {"repo": "Econogoat/Krypt0_LORA", "trigger": "Krypt0", "adapter_name": "krypt0"}
-LLM_MODEL_ID = "google/gemma-3n-E4B-it"
-SYSTEM_PROMPT = """You are a creative assistant that enhances user prompts for an AI image generation model.
-Your task is to take a user's simple idea and expand it into a rich, detailed, and visually descriptive prompt.
-Focus on cinematic lighting, intricate details, atmosphere, and a strong artistic style.
-Do NOT add the trigger word 'Krypt0', it will be added automatically later.
-Reply ONLY with the enhanced prompt, without any introduction or explanation."""
+# --- Main Configuration ---
+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 ---
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+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.")
+
+# 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']
+)
+print("LoRA loaded successfully.")
+
 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 calculate_dimensions(aspect_ratio, resolution):
+    """Calculates width and height based on aspect ratio and base resolution."""
     resolution = int(resolution)
-    if aspect_ratio == "Square (1:1)": width, height = resolution, resolution
-    elif aspect_ratio == "Portrait (9:16)": width, height = int(resolution * 9 / 16), resolution
-    elif aspect_ratio == "Landscape (16:9)": width, height = resolution, int(resolution * 9 / 16)
-    elif aspect_ratio == "Ultrawide (21:9)": width, height = resolution, int(resolution * 9 / 21)
-    else: width, height = resolution, resolution
+    if aspect_ratio == "Square (1:1)":
+        width, height = resolution, resolution
+    elif aspect_ratio == "Portrait (9:16)":
+        width, height = int(resolution * 9 / 16), resolution
+    elif aspect_ratio == "Landscape (16:9)":
+        width, height = resolution, int(resolution * 9 / 16)
+    elif aspect_ratio == "Ultrawide (21:9)":
+        width, height = resolution, int(resolution * 9 / 21)
+    else: # Fallback
+        width, height = resolution, resolution
+    # Ensure dimensions are multiples of 64 for optimal performance
     width = (width // 64) * 64
     height = (height // 64) * 64
     return width, height
 
-def update_history(new_image, history):
-    if new_image is None: return history
-    if history is None: history = []
-    history.insert(0, new_image)
-    return history
-
-@spaces.GPU(duration=180)
-def run_generation(prompt, lora_scale, cfg_scale, steps, randomize_seed, seed, aspect_ratio, base_resolution, 
-                   # On reçoit l'état actuel des modèles en entrée
-                   state_pipe, state_llm_model, state_llm_processor, state_good_vae, 
-                   progress=gr.Progress(track_tqdm=True)):
-    
-    # --- CHARGEMENT À LA VOLÉE AU PREMIER CLIC, EN UTILISANT gr.State ---
-    # La condition est maintenant basée sur l'état passé en argument, pas sur une variable globale fragile.
-    if state_pipe is None:
-        gr.Info("First run: Loading all models... This will take a moment.")
-        print("First run: Loading all models inside GPU context...")
-        
-        device = "cuda"
-        dtype = torch.bfloat16
-
-        print("Loading LLM...")
-        state_llm_processor = AutoProcessor.from_pretrained(LLM_MODEL_ID, token=HF_TOKEN)
-        state_llm_model = AutoModelForImageTextToText.from_pretrained(LLM_MODEL_ID, torch_dtype=dtype, token=HF_TOKEN)
-        
-        print("Loading diffusion models...")
-        taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
-        state_good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", token=HF_TOKEN, torch_dtype=dtype)
-        state_pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", vae=taef1, token=HF_TOKEN, torch_dtype=dtype)
-        
-        print("Loading LoRA...")
-        state_pipe.load_lora_weights(KRYPTO_LORA['repo'], low_cpu_mem_usage=False, adapter_name=KRYPTO_LORA['adapter_name'])
-        
-        state_pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(state_pipe)
-        print("All models loaded and stored in session state.")
-    
-    # --- DÉBUT DU PROCESSUS NORMAL ---
-    if not prompt: raise gr.Error("Prompt cannot be empty.")
-    
-    device = "cuda"
-    dtype = torch.bfloat16
-    
-    # --- 1. Amélioration du prompt avec le LLM ---
-    gr.Info("Enhancing prompt with LLM...")
-    state_llm_model.to(device)
-    messages = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt}]
-    inputs = state_llm_processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_tensors="pt").to(device)
-    
-    with torch.inference_mode():
-        outputs = state_llm_model.generate(**inputs, max_new_tokens=150)
-    
-    enhanced_prompt = state_llm_processor.batch_decode(outputs, skip_special_tokens=True)[0].split("assistant\n")[-1].strip()
-    state_llm_model.to("cpu"); torch.cuda.empty_cache()
-
-    # --- 2. Génération d'image ---
-    gr.Info("Prompt enhanced. Starting image generation...")
-    prompt_mash = f"{KRYPTO_LORA['trigger']}, {enhanced_prompt}"
-    state_pipe.set_adapters([KRYPTO_LORA['adapter_name']], adapter_weights=[lora_scale])
-    if randomize_seed: seed = random.randint(0, MAX_SEED)
-    width, height = calculate_dimensions(aspect_ratio, base_resolution)
-    
-    state_pipe.to(device); state_good_vae.to(device)
+def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress):
+    """Generator function for text-to-image with live preview."""
+    # The parent @spaces.GPU function has already allocated a GPU
+    pipe.to(device) 
     generator = torch.Generator(device=device).manual_seed(seed)
     
-    image_generator = state_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, joint_attention_kwargs={"scale": 1.0},
-        output_type="pil", good_vae=state_good_vae
+    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,
+        joint_attention_kwargs={"scale": 1.0},
+        output_type="pil",
+        good_vae=good_vae,
     )
+    # Yield previews and the final image
     final_image = None
     for i, image in enumerate(image_generator):
         final_image = image
         progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {i + 1}; --total: {steps};"></div></div>'
-        # On retourne les états inchangés pendant la prévisualisation
-        yield image, seed, gr.update(value=progress_bar, visible=True), state_pipe, state_llm_model, state_llm_processor, state_good_vae
+        yield image, gr.update(value=progress_bar, visible=True)
+    yield final_image, gr.update(visible=False)
+
+def update_history(new_image, history):
+    """Adds the new image to the history gallery."""
+    if new_image is None: # Don't add empty images on error
+        return history
+    if history is None:
+        history = []
+    history.insert(0, new_image)
+    return history
     
-    state_pipe.to("cpu"); state_good_vae.to("cpu")
+@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)):
+    if not prompt:
+        raise gr.Error("Prompt cannot be empty.")
+
+    prompt_mash = f"{KRYPTO_LORA['trigger']}, {prompt}"
+    print("Final prompt:", prompt_mash)
+    
+    # Activate the LoRA adapter with the slider's weight
+    pipe.set_adapters([KRYPTO_LORA['adapter_name']], adapter_weights=[lora_scale])
+    print(f"Adapter '{KRYPTO_LORA['adapter_name']}' activated with weight {lora_scale}.")
     
-    # On retourne l'image finale ET l'état mis à jour des modèles pour qu'ils soient conservés pour le prochain appel
-    yield final_image, seed, gr.update(visible=False), state_pipe, state_llm_model, state_llm_processor, state_good_vae
+    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.")
+
+    # The function now only handles text-to-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
 
-# --- UI ---
-css = ''' #title{text-align: center} #title h1{font-size: 3em; display:inline-flex; align-items:center} #title img{width: 80px; margin-right: 0.25em} .progress-container {width: 100%;height: 30px;background-color: #f0f0f0;border-radius: 15px;overflow: hidden;margin-bottom: 20px} .progress-bar {height: 100%;background-color: #4f46e5;width: calc(var(--current) / var(--total) * 100%);transition: width 0.1s ease-in-out} #random_prompt_btn{max-width: 2.5em; min-width: 2.5em !important; height: 100% !important;} '''
-with gr.Blocks(css=css, theme=gr.themes.Soft()) as app:
-    # Déclaration des états qui contiendront nos modèles
-    state_pipe = gr.State(None)
-    state_llm_model = gr.State(None)
-    state_llm_processor = gr.State(None)
-    state_good_vae = gr.State(None)
+# --- User Interface (Gradio) ---
+css = '''
+#title{text-align: center}
+#title h1{font-size: 3em; display:inline-flex; align-items:center}
+#title img{width: 80px; margin-right: 0.25em}
+.progress-container {width: 100%;height: 30px;background-color: #f0f0f0;border-radius: 15px;overflow: hidden;margin-bottom: 20px}
+.progress-bar {height: 100%;background-color: #4f46e5;width: calc(var(--current) / var(--total) * 100%);transition: width 0.1s ease-in-out}
+#random_prompt_btn{max-width: 2.5em; min-width: 2.5em !important; height: 100% !important;}
+'''
 
-    gr.HTML("""<div id='title'><h1><img src="https://huggingface.co/Econogoat/KRYPTO_LORA/resolve/main/krypt0.png" alt="LoRA"> Krypt0 Image Generator</h1><br><span>Generate images with the Krypt0 artistic style</span></div>""")
+with gr.Blocks(css=css, theme=gr.themes.Soft()) as app:
+    # --- Header ---
+    with gr.Row():
+        gr.HTML(
+            """<div id='title'><h1><img src="https://huggingface.co/Econogoat/KRYPTO_LORA/resolve/main/krypt0.png" alt="LoRA"> Krypt0 Image Generator</h1><br><span>Generate images with the Krypt0 artistic style</span></div>"""
+        )
+    
     with gr.Row():
+        # --- LEFT COLUMN: CONTROLS ---
         with gr.Column(scale=3):
+            # Prompt and Style Controls
             with gr.Group():
                 with gr.Row():
                     random_prompt_btn = gr.Button("🎲", elem_id="random_prompt_btn")
                     prompt = gr.Textbox(label="Prompt", lines=2, placeholder="e.g., a portrait of a warrior queen", scale=8)
+                
                 lora_scale = gr.Slider(label="Krypt0 Style Strength", minimum=0, maximum=2, step=0.05, value=0.9)
-            aspect_ratio = gr.Radio(label="Aspect Ratio", choices=["Square (1:1)", "Portrait (9:16)", "Landscape (16:9)", "Ultrawide (21:9)"], value="Square (1:1)")
+            
+            # Image Shape Controls
+            aspect_ratio = gr.Radio(
+                label="Aspect Ratio",
+                choices=["Square (1:1)", "Portrait (9:16)", "Landscape (16:9)", "Ultrawide (21:9)"],
+                value="Square (1:1)"
+            )
+            
+            # Advanced Settings
             with gr.Accordion("Advanced Settings", open=True):
                 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)
                 cfg_scale = gr.Slider(label="Guidance (CFG Scale)", minimum=1, maximum=10, step=0.5, value=3.5)
+                
                 with gr.Row():
                     randomize_seed = gr.Checkbox(True, label="Random Seed")
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+            
             generate_button = gr.Button("Generate", variant="primary")
+            
+        # --- RIGHT COLUMN: RESULTS ---
         with gr.Column(scale=2):
             progress_bar = gr.Markdown(elem_id="progress", visible=False)
             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])
-    
-    # On ajoute les états aux inputs et outputs de l'événement de génération
-    # C'est la boucle qui assure la persistence des modèles
-    model_states = [state_pipe, state_llm_model, state_llm_processor, state_good_vae]
-    
+    # --- Event Logic ---
+    def get_random_prompt():
+        return random.choice(prompt_values)
+        
+    random_prompt_btn.click(
+        fn=get_random_prompt,
+        inputs=[],
+        outputs=[prompt]
+    )
+
     generation_event = gr.on(
-        triggers=[generate_button.click, prompt.submit], 
-        fn=run_generation, 
-        inputs=[prompt, lora_scale, cfg_scale, steps, randomize_seed, seed, aspect_ratio, base_resolution] + model_states, 
-        outputs=[result, seed, progress_bar] + model_states
+        triggers=[generate_button.click, prompt.submit],
+        fn=run_generation,
+        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,
     )
-    generation_event.then(fn=update_history, inputs=[result, history_gallery], outputs=history_gallery)
 
 app.queue(max_size=20)
 app.launch()