Fix lora mem issue with large res

app.py

@@ -94,6 +94,7 @@ MOD_VALUE = 8
 DEFAULT_H_SLIDER_VALUE = 512
 DEFAULT_W_SLIDER_VALUE = 768
 NEW_FORMULA_MAX_AREA = 768.0 * 512.0
+LORA_MAX_AREA = 640.0 * 480.0 # Max pixel area when a LoRA is active
 
 SLIDER_MIN_H, SLIDER_MAX_H = 128, 896
 SLIDER_MIN_W, SLIDER_MAX_W = 128, 896
@@ -126,43 +127,56 @@ def parse_lset_prompt(lset_prompt):
         
     return resolved_prompt
 
-def handle_lora_selection_change(preset_name, current_prompt):
+def handle_lora_selection_change(preset_name, current_prompt, current_h, current_w, aspect_ratio):
     """
     When a preset is selected, this function finds the corresponding .lset file,
-    parses it, and appends the generated prompt to the current prompt.
+    parses it, appends the prompt, and resizes dimensions if they are too large.
     """
+    # Initialize updates to avoid changing UI elements unnecessarily
+    prompt_update = gr.update()
+    h_update = gr.update()
+    w_update = gr.update()
+
     if not preset_name or preset_name == "None":
-        return gr.update() # No preset selected, do not change the prompt.
+        return prompt_update, h_update, w_update
 
+    # --- Handle Prompt ---
     try:
-        # The preset_name is the filename without extension
         lset_filename = f"{preset_name}.lset"
-        
-        # Download the .lset file from the same subfolder as the LoRA
         lset_path = hf_hub_download(
-            repo_id=I2V_LORA_REPO_ID,
-            filename=lset_filename,
-            subfolder=I2V_LORA_SUBFOLDER,
-            repo_type='model'
+            repo_id=I2V_LORA_REPO_ID, filename=lset_filename,
+            subfolder=I2V_LORA_SUBFOLDER, repo_type='model'
         )
-        
         with open(lset_path, 'r', encoding='utf-8') as f:
             lset_data = json.load(f)
         
-        lset_prompt_raw = lset_data.get("prompt")
-        if not lset_prompt_raw:
-            return gr.update()
-
-        resolved_prompt = parse_lset_prompt(lset_prompt_raw)
-        # Append with newlines for separation. If current prompt is empty, strip() handles it.
-        new_prompt = f"{current_prompt}\n\n{resolved_prompt}".strip()
-        gr.Info(f"✅ Appended prompt from '{lset_filename}'. Replace highlighted text like __this__.")
-        return gr.update(value=new_prompt)
+        if lset_prompt_raw := lset_data.get("prompt"):
+            resolved_prompt = parse_lset_prompt(lset_prompt_raw)
+            new_prompt = f"{current_prompt}\n\n{resolved_prompt}".strip()
+            gr.Info(f"✅ Appended prompt from '{lset_filename}'. Replace highlighted text like __this__.")
+            prompt_update = gr.update(value=new_prompt)
     except Exception as e:
-        # This should be less common now, but good to keep for safety.
         print(f"Info: Could not process .lset for '{preset_name}'. Reason: {e}")
         gr.Info(f"ℹ️ Error processing preset '{preset_name}'.")
-        return gr.update()
+
+    # --- Handle Resolution ---
+    if current_h * current_w > LORA_MAX_AREA:
+        gr.Info(f"Resolution too high for LoRA. Scaling down to a 640x480 equivalent area.")
+        # aspect_ratio is W/H
+        if aspect_ratio > 0:
+            calc_w = round(np.sqrt(LORA_MAX_AREA * aspect_ratio))
+            calc_h = round(np.sqrt(LORA_MAX_AREA / aspect_ratio))
+            
+            new_h = max(MOD_VALUE, (calc_h // MOD_VALUE) * MOD_VALUE)
+            new_w = max(MOD_VALUE, (calc_w // MOD_VALUE) * MOD_VALUE)
+            
+            h_update = gr.update(value=new_h)
+            w_update = gr.update(value=new_w)
+        else: # Fallback if aspect ratio is invalid
+            h_update = gr.update(value=480)
+            w_update = gr.update(value=640)
+
+    return prompt_update, h_update, w_update
 
 # --- Helper Functions ---
 def sanitize_prompt_for_filename(prompt: str, max_len: int = 60) -> str:
@@ -259,6 +273,21 @@ def generate_i2v_video(input_image, prompt, height, width,
     target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
     target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)
     
+    # If a LoRA is used, enforce max resolution as a safety net
+    if preset_name and preset_name != "None":
+        if target_h * target_w > LORA_MAX_AREA:
+            print(f"⚠️ Warning: Resolution {target_w}x{target_h} is too high for LoRA. Rescaling to fit max area.")
+            aspect_ratio = target_w / target_h if target_h > 0 else 1.0
+            
+            # Re-calculate w and h based on max area
+            target_w = round(np.sqrt(LORA_MAX_AREA * aspect_ratio))
+            target_h = round(np.sqrt(LORA_MAX_AREA / aspect_ratio))
+
+            # Snap to MOD_VALUE
+            target_h = max(MOD_VALUE, (target_h // MOD_VALUE) * MOD_VALUE)
+            target_w = max(MOD_VALUE, (target_w // MOD_VALUE) * MOD_VALUE)
+            print(f"   - Rescaled to: {target_w}x{target_h}")
+
     # Calculate and adjust num_frames to be compatible with video codecs
     target_frames = int(round(duration_seconds * FIXED_FPS))
     adjusted_frames = 4 * round((target_frames - 1) / 4) + 1
@@ -393,8 +422,8 @@ with gr.Blocks() as demo:
     # I2V Handlers
     i2v_preset_name.change(
         fn=handle_lora_selection_change,
-        inputs=[i2v_preset_name, i2v_prompt],
-        outputs=[i2v_prompt]
+        inputs=[i2v_preset_name, i2v_prompt, i2v_height, i2v_width, i2v_aspect_ratio],
+        outputs=[i2v_prompt, i2v_height, i2v_width]
     )
     i2v_input_image.upload(
         fn=handle_image_upload_for_dims_wan,