Big code cleanup

app.py

@@ -139,20 +139,19 @@ def parse_lset_prompt(lset_prompt):
         resolved_prompt = resolved_prompt.replace(f"{{{key}}}", highlighted_value)
     return resolved_prompt
 
-def handle_lora_selection_change(preset_name: str, base_prompt: str):
+def handle_lora_selection_change(preset_name: str, current_prompt: str):
     """
-    When a preset is selected, this function combines the base_prompt (from state)
-    with the new LoRA's trigger words. This makes the update idempotent.
+    Appends the selected LoRA's trigger words to the current prompt text
+    and controls the visibility of the weight slider.
     """
-    # Default state for the slider is hidden and non-interactive
-    lora_slider_update = gr.update(visible=False, interactive=False)
-    
-    # If "None" is selected, the displayed prompt is just the base prompt.
+    # If "None" is selected, do nothing to the prompt and hide the slider.
     if not preset_name or preset_name == "None":
         gr.Info("LoRA cleared.")
-        return gr.update(value=base_prompt), lora_slider_update
+        # Return the prompt unchanged, and hide the slider.
+        return gr.update(value=current_prompt), gr.update(visible=False, interactive=False)
 
     try:
+        # Fetch the trigger words from the LoRA's .lset file.
         lset_filename = f"{preset_name}.lset"
         lset_path = hf_hub_download(
             repo_id=DYNAMIC_LORA_REPO_ID,
@@ -166,38 +165,26 @@ def handle_lora_selection_change(preset_name: str, base_prompt: str):
             lset_data = json.loads(lset_content)
             lset_prompt_raw = lset_data.get("prompt")
         except json.JSONDecodeError:
-            print(f"Info: '{lset_filename}' is not JSON. Treating as plain text prompt.")
             lset_prompt_raw = lset_content
 
         if lset_prompt_raw:
-            resolved_prompt = parse_lset_prompt(lset_prompt_raw)
-            separator = ", " if base_prompt and not base_prompt.endswith((",", " ")) else ""
-            # The new prompt is always constructed from the base prompt and the new triggers.
-            new_prompt = f"{base_prompt}{separator}{resolved_prompt}".strip()
-            gr.Info(f"✅ Appended triggers from '{preset_name}'. Replace highlighted text like __this__.")
-            
-            # On success, update the prompt and make the slider visible and interactive
-            new_prompt_update = gr.update(value=new_prompt)
-            lora_slider_update = gr.update(visible=True, interactive=True)
-            return new_prompt_update, lora_slider_update
+            # Append the new trigger words to the current prompt.
+            trigger_words = parse_lset_prompt(lset_prompt_raw)
+            separator = ", " if current_prompt and not current_prompt.endswith((",", " ")) else ""
+            new_prompt = f"{current_prompt}{separator}{trigger_words}".strip()
+            gr.Info(f"✅ Appended triggers from '{preset_name}'. You can now edit them.")
+            # Return the updated prompt and show the slider.
+            return gr.update(value=new_prompt), gr.update(visible=True, interactive=True)
         else:
+            # If the .lset file has no prompt, do nothing.
             gr.Info(f"ℹ️ No prompt found in '{preset_name}.lset'. Prompt unchanged.")
-            # If no triggers, the prompt is just the base prompt.
-            return gr.update(value=base_prompt), lora_slider_update
+            return gr.update(value=current_prompt), gr.update(visible=True, interactive=True)
 
     except Exception as e:
         print(f"Info: Could not process .lset for '{preset_name}'. Reason: {e}")
         gr.Warning(f"⚠️ Could not load triggers for '{preset_name}'.")
-        # On error, revert to just the base prompt.
-        return gr.update(value=base_prompt), lora_slider_update
-
-def set_base_prompt(current_prompt_text):
-    """
-    Called when the user clicks the 'Set as Base Prompt' button.
-    This updates the base prompt state and resets the LoRA selection.
-    """
-    gr.Info("New base prompt set. You can now select a LoRA to add triggers.")
-    return current_prompt_text, "None", gr.update(visible=False, interactive=False)
+        # On error, return the prompt unchanged but still show the slider.
+        return gr.update(value=current_prompt), gr.update(visible=True, interactive=True)
 
 
 def _manage_lora_state(pipe, selected_lora: str, lora_weight: float) -> bool:
@@ -208,16 +195,13 @@ def _manage_lora_state(pipe, selected_lora: str, lora_weight: float) -> bool:
         bool: True if a dynamic LoRA was loaded, False otherwise.
     """
     # Pre-emptive cleanup of any previously loaded dynamic adapter.
-    # This is more robust than relying only on the `finally` block of the previous run.
     try:
-        # --- FIX: Use delete_adapters to remove a specific adapter by name ---
         pipe.delete_adapters([DYNAMIC_LORA_ADAPTER_NAME])
         print("🧼 Pre-emptively unloaded previous dynamic LoRA.")
     except Exception:
         pass  # No dynamic lora was present, which is a clean state.
 
     if not selected_lora or selected_lora == "None":
-        # This run uses no dynamic LoRA. Ensure only the base LoRA is active.
         pipe.set_adapters([FUSIONX_ADAPTER_NAME], adapter_weights=[FUSIONX_LORA_WEIGHT])
         print("ℹ️ No dynamic LoRA selected. Using base LoRA only.")
         return False
@@ -226,7 +210,6 @@ def _manage_lora_state(pipe, selected_lora: str, lora_weight: float) -> bool:
     print(f"🚀 Processing preset: {selected_lora} with weight {lora_weight}")
     lora_filename = None
     try:
-        # Attempt to get the real LoRA filename from the .lset file
         lset_filename = f"{selected_lora}.lset"
         lset_path = hf_hub_download(
             repo_id=DYNAMIC_LORA_REPO_ID,
@@ -268,27 +251,18 @@ def load_pipelines():
 
     print("\n🚀 Loading T2V pipeline with base LoRA...")
     try:
-        # To avoid potential cache duplication and storage issues, we load the
-        # pipeline directly, then replace the VAE with a float32 version for stability.
         t2v_pipe = DiffusionPipeline.from_pretrained(
             T2V_BASE_MODEL_ID,
-            torch_dtype=torch.bfloat16, # Load in bfloat16 to save memory
+            torch_dtype=torch.bfloat16,
         )
         print("✅ Base pipeline loaded. Overriding VAE with float32 version...")
-
-        # The VAE often works better in float32. We reload it and replace it in the pipeline.
-        # Using the specific AutoencoderKLWan class is more robust than the generic AutoModel.
         vae_fp32 = AutoencoderKLWan.from_pretrained(T2V_BASE_MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
         t2v_pipe.vae = vae_fp32
-
         t2v_pipe.to("cuda")
         print("✅ Pipeline configured. Loading and activating base FusionX LoRA...")
-
-        # Load and set the base LoRA that is always active
         t2v_pipe.load_lora_weights(FUSIONX_LORA_REPO, weight_name=FUSIONX_LORA_FILE, adapter_name=FUSIONX_ADAPTER_NAME)
         t2v_pipe.set_adapters([FUSIONX_ADAPTER_NAME], adapter_weights=[FUSIONX_LORA_WEIGHT])
         print("✅ T2V pipeline with base LoRA is ready.")
-
     except Exception as e:
         print(f"❌ CRITICAL ERROR: Failed to load T2V pipeline. T2V will be disabled. Reason: {e}")
         traceback.print_exc()
@@ -296,8 +270,6 @@ def load_pipelines():
 
     print("\n🤖 Loading LLM for Prompt Enhancement...")
     try:
-        # In a ZeroGPU environment, we must load models on the CPU at startup.
-        # The model will be moved to the GPU inside the decorated function.
         enhancer_pipe = pipeline("text-generation", model=ENHANCER_MODEL_ID, torch_dtype=torch.bfloat16, device="cpu")
         print("✅ LLM Prompt Enhancer loaded successfully (on CPU).")
     except Exception as e:
@@ -313,17 +285,11 @@ def load_pipelines():
 def enhance_prompt_with_llm(prompt: str, enhancer_pipeline):
     """
     Uses the loaded LLM to enhance a given prompt.
-    This function manually handles tensor placement to avoid device mismatches in a ZeroGPU environment.
     """
     if enhancer_pipeline is None:
-        print("LLM enhancer not available, returning original prompt.")
         gr.Warning("LLM enhancer is not available.")
-        return prompt
+        return prompt, gr.update(), gr.update()
 
-    # In a Hugging Face ZeroGPU Space, the GPU is provisioned on-demand for functions
-    # decorated with @spaces.GPU and de-provisioned afterward. Therefore, the model,
-    # which is loaded on the CPU at startup, must be moved to the GPU for every call.
-    # The "Moving enhancer model to CUDA..." message is expected and correct for this setup.
     if enhancer_pipeline.model.device.type != 'cuda':
         print("Moving enhancer model to CUDA for on-demand GPU execution...")
         enhancer_pipeline.model.to("cuda")
@@ -332,64 +298,31 @@ def enhance_prompt_with_llm(prompt: str, enhancer_pipeline):
     print(f"Enhancing prompt: '{prompt}'")
 
     try:
-        # 1. Get the tokenizer from the pipeline.
         tokenizer = enhancer_pipeline.tokenizer
-
-        # FIX: Set pad_token to eos_token if not set. This is a common requirement for
-        # models like Qwen2 and helps prevent warnings about attention masks.
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
-
-        # 2. Apply the chat template and tokenize. This returns a dictionary containing
-        #    'input_ids' and 'attention_mask' as PyTorch tensors.
-        tokenized_inputs = tokenizer.apply_chat_template(
-            messages,
-            tokenize=True,
-            add_generation_prompt=True,
-            return_tensors="pt"
-        )
+        tokenized_inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
         
-        # 3. FIX: Move each tensor in the dictionary to the CUDA device.
-        # 3. FIX: The tokenizer might return a single tensor instead of a dictionary.
-        # We handle both cases to make the code more robust.
         if isinstance(tokenized_inputs, torch.Tensor):
-            # If we get a single tensor, assume it's input_ids
             inputs_on_cuda = {"input_ids": tokenized_inputs.to("cuda")}
-            # Manually create the attention mask as it's good practice for generate()
             inputs_on_cuda["attention_mask"] = torch.ones_like(inputs_on_cuda["input_ids"])
         else:
-            # If we get a dictionary, move all its tensors to cuda
             inputs_on_cuda = {k: v.to("cuda") for k, v in tokenized_inputs.items()}
 
-        # 4. Use the model's generate() method, unpacking the dictionary to pass
-        #    both `input_ids` and `attention_mask`. This resolves the warning.
-        generated_ids = enhancer_pipeline.model.generate(
-            **inputs_on_cuda,
-            max_new_tokens=256,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.95
-        )
-
-        # 5. The output from generate() includes the input tokens. We need to decode only the newly generated part.
+        generated_ids = enhancer_pipeline.model.generate(**inputs_on_cuda, max_new_tokens=256, do_sample=True, temperature=0.7, top_p=0.95)
         input_token_length = inputs_on_cuda['input_ids'].shape[1]
         newly_generated_ids = generated_ids[:, input_token_length:]
-        
-        # 6. Decode the new tokens back into a string.
         final_answer = tokenizer.decode(newly_generated_ids[0], skip_special_tokens=True)
 
         print(f"Enhanced prompt: '{final_answer.strip()}'")
-        # Return the enhanced prompt and also reset the LoRA dropdown
+        # The enhanced prompt overwrites the textbox. The LoRA selection is reset.
         return final_answer.strip(), "None", gr.update(visible=False, interactive=False)
     except Exception as e:
         print(f"❌ Error during prompt enhancement: {e}")
-        # Adding full traceback for better debugging in the console
         traceback.print_exc()
         gr.Warning(f"An error occurred during prompt enhancement. See console for details.")
-        return prompt, gr.update(), gr.update() # Return original prompt, don't change LoRA
+        return prompt, gr.update(), gr.update()
     finally:
-        # Explicitly empty the CUDA cache to help release GPU memory.
-        # This can help resolve intermittent issues where the GPU remains active.
         print("🧹 Clearing CUDA cache after prompt enhancement...")
         torch.cuda.empty_cache()
 
@@ -408,11 +341,12 @@ def generate_t2v_video(
     if not prompt:
         raise gr.Error("Please enter a prompt for Text-to-Video generation.")
 
+    # --- The prompt from the textbox is now the final prompt. No more combining. ---
+    final_prompt = prompt
+
     target_h = max(MOD_VALUE, (height // MOD_VALUE) * MOD_VALUE)
     target_w = max(MOD_VALUE, (width // MOD_VALUE) * MOD_VALUE)
 
-    # Calculate the initial number of frames based on duration and model constraints.
-    # The model requires (num_frames - 1) to be divisible by 4.
     requested_frames = int(round(duration_seconds * T2V_FIXED_FPS))
     frames_minus_one = requested_frames - 1
     valid_frames_minus_one = round(frames_minus_one / 4.0) * 4
@@ -425,28 +359,23 @@ def generate_t2v_video(
     lora_loaded = False
 
     try:
-        # Manage the LoRA state and get a flag indicating if a dynamic LoRA was loaded.
-        lora_loaded = _manage_lora_state(
-            pipe=t2v_pipe,
-            selected_lora=selected_lora,
-            lora_weight=lora_weight
-        )
+        lora_loaded = _manage_lora_state(pipe=t2v_pipe, selected_lora=selected_lora, lora_weight=lora_weight)
 
         print("\n--- Starting T2V Generation ---")
-        print(f"Prompt: {prompt}")
+        print(f"Final Prompt: {final_prompt}")
         print(f"Resolution: {target_w}x{target_h}, Frames: {num_frames}, Seed: {current_seed}")
         print(f"Steps: {steps}, Guidance: 1.0 (fixed for FusionX)")
         print("---------------------------------")
 
         with torch.inference_mode():
             output_frames_list = t2v_pipe(
-                prompt=prompt, negative_prompt=negative_prompt,
+                prompt=final_prompt, negative_prompt=negative_prompt,
                 height=target_h, width=target_w, num_frames=num_frames,
                 guidance_scale=1.0, num_inference_steps=int(steps),
                 generator=torch.Generator(device="cuda").manual_seed(current_seed)
             ).frames[0]
 
-        sanitized_prompt = sanitize_prompt_for_filename(prompt)
+        sanitized_prompt = sanitize_prompt_for_filename(final_prompt)
         filename = f"t2v_{sanitized_prompt}_{current_seed}.mp4"
         temp_dir = tempfile.mkdtemp()
         video_path = os.path.join(temp_dir, filename)
@@ -457,28 +386,20 @@ def generate_t2v_video(
         return video_path, current_seed, gr.File(value=video_path, visible=True, label=download_label)
     
     except Exception as e:
-        # Broad exception to catch any error during generation and ensure cleanup still happens
         print(f"❌ An error occurred during video generation: {e}")
         traceback.print_exc()
         raise gr.Error("Video generation failed. Please check the logs for details.")
 
     finally:
-        # --- CLEANUP ---
-        # This block ensures the dynamic LoRA is removed after every run,
-        # resetting the pipeline to a clean state for the next user.
         if lora_loaded:
             print(f"🧼 Cleaning up dynamic LoRA: {selected_lora}")
             try:
-                # --- FIX: Use delete_adapters to correctly remove the adapter ---
                 t2v_pipe.delete_adapters([DYNAMIC_LORA_ADAPTER_NAME])
-                # IMPORTANT: Reset adapters back to just the base LoRA for the next run.
                 t2v_pipe.set_adapters([FUSIONX_ADAPTER_NAME], adapter_weights=[FUSIONX_LORA_WEIGHT])
                 print("✅ Cleanup complete. Pipeline reset to base LoRA state.")
             except Exception as e:
                 print(f"⚠️ Error during LoRA cleanup: {e}. State may be inconsistent.")
         
-        # Explicitly empty the CUDA cache to help release GPU memory.
-        # This can help resolve intermittent issues where the GPU remains active.
         print("🧹 Clearing CUDA cache after video generation...")
         torch.cuda.empty_cache()
 
@@ -488,11 +409,8 @@ def generate_t2v_video(
 def build_ui(t2v_pipe, enhancer_pipe, available_loras):
     """Creates and configures the Gradio UI."""
     with gr.Blocks(theme=gr.themes.Soft(), css=".main-container { max-width: 1080px; margin: auto; }") as demo:
-        # --- Add a state component to reliably store the user's base prompt ---
-        base_prompt_state = gr.State(value=DEFAULT_PROMPT_T2V)
-        
         gr.Markdown("# ✨ Wan 2.1 Text-to-Video Suite with Dynamic LoRAs")
-        gr.Markdown("Generate videos from text, enhanced by the base `FusionX` LoRA and your choice of dynamic style LoRA.")
+        gr.Markdown("Generate videos from text. Edit the prompt below. Selecting a LoRA will append its triggers to your prompt.")
 
         with gr.Tabs():
             with gr.TabItem("✍️ Text-to-Video", id="t2v_tab", interactive=t2v_pipe is not None):
@@ -505,22 +423,17 @@ def build_ui(t2v_pipe, enhancer_pipe, available_loras):
                                 label="✏️ Prompt", value=DEFAULT_PROMPT_T2V, lines=4,
                                 placeholder="e.g., A cinematic drone shot flying over a futuristic city at night..."
                             )
-                            with gr.Row():
-                                t2v_enhance_btn = gr.Button(
-                                    "🤖 Enhance Prompt with AI",
-                                    interactive=enhancer_pipe is not None
-                                )
-                                # --- FIX: Add a button to explicitly set the base prompt ---
-                                set_base_btn = gr.Button(
-                                    "📌 Set as Base Prompt"
-                                )
+                            t2v_enhance_btn = gr.Button(
+                                "🤖 Enhance Prompt with AI",
+                                interactive=enhancer_pipe is not None
+                            )
 
                             with gr.Group():
                                 t2v_lora_preset = gr.Dropdown(
                                     label="🎨 Dynamic Style LoRA (Optional)",
                                     choices=available_loras,
                                     value="None",
-                                    info="Adds a secondary style LoRA. Replaces previous LoRA triggers."
+                                    info="Appends style triggers to the prompt text above."
                                 )
                                 t2v_lora_weight = gr.Slider(
                                     label="💪 LoRA Weight", minimum=0.0, maximum=2.0, step=0.05, value=0.8, 
@@ -549,40 +462,24 @@ def build_ui(t2v_pipe, enhancer_pipe, available_loras):
                             t2v_download = gr.File(label="📥 Download Video", visible=False)
 
         if t2v_pipe is not None:
-            # Create a partial function that has the enhancer_pipe "baked in".
             enhance_fn = partial(enhance_prompt_with_llm, enhancer_pipeline=enhancer_pipe)
             
-            # --- Wire up the new state-based event handlers ---
-            
-            # 1. When user clicks the "Set as Base" button after a manual edit.
-            set_base_btn.click(
-                fn=set_base_prompt,
-                inputs=[t2v_prompt],
-                outputs=[base_prompt_state, t2v_lora_preset, t2v_lora_weight],
-                queue=False 
-            )
-            
-            # 2. When the user enhances the prompt with the LLM. This also creates a new base prompt.
+            # 1. When the user enhances the prompt with the LLM.
             t2v_enhance_btn.click(
                 fn=enhance_fn,
                 inputs=[t2v_prompt],
-                # The enhance function now also resets the LoRA dropdown and slider
                 outputs=[t2v_prompt, t2v_lora_preset, t2v_lora_weight]
-            ).then(
-                fn=lambda p: p, # A simple function to pass the new prompt through
-                inputs=[t2v_prompt],
-                outputs=[base_prompt_state] # Update the base prompt state with the enhanced version
             )
 
-            # 3. When the user selects a LoRA from the dropdown.
+            # 2. When the user selects a LoRA from the dropdown.
             t2v_lora_preset.change(
                 fn=handle_lora_selection_change,
-                # The input is now the reliable base_prompt_state, not the textbox
-                inputs=[t2v_lora_preset, base_prompt_state],
+                # Pass the current prompt text in, get the new text back out.
+                inputs=[t2v_lora_preset, t2v_prompt],
                 outputs=[t2v_prompt, t2v_lora_weight]
             )
             
-            # 4. When the user clicks the final generate button.
+            # 3. When the user clicks the final generate button.
             t2v_generate_btn.click(
                 fn=generate_t2v_video,
                 inputs=[
@@ -599,7 +496,6 @@ def build_ui(t2v_pipe, enhancer_pipe, available_loras):
 if __name__ == "__main__":
     t2v_pipe, enhancer_pipe = load_pipelines()
     
-    # Fetch LoRAs only if the main pipeline loaded successfully
     available_loras = []
     if t2v_pipe:
         available_loras = get_available_presets(DYNAMIC_LORA_REPO_ID, DYNAMIC_LORA_SUBFOLDER)