Update app.py

app.py

@@ -195,24 +195,64 @@ def _list_loras():
     return [LORA_NONE]
 
 
-def _add_lora_to_workflow(workflow, model_node_ref, lora_name, lora_strength):
-    """Inserts a LoraLoaderModelOnly node between the diffusion model and the
-    sampler when a LoRA is selected. Uses *ModelOnly* on purpose: the text
-    encoder here (qwen3vl) is loaded separately via CLIPLoader and isn't part
-    of the UNET checkpoint, so a regular LoraLoader (which also patches CLIP)
-    would be reaching for a CLIP input that doesn't exist in this graph."""
-    if not lora_name or lora_name == LORA_NONE:
-        return model_node_ref
-
-    workflow["20"] = {
-        "class_type": "LoraLoaderModelOnly",
-        "inputs": {
-            "model": model_node_ref,
-            "lora_name": lora_name,
-            "strength_model": float(lora_strength),
-        },
-    }
-    return ["20", 0]
+def _download_dynamic_lora(repo, filename):
+    """Downloads a dynamic LoRA from HuggingFace on-the-fly inside the execution flow.
+    Returns the relative path from models/loras for ComfyUI loader."""
+    if not repo or not filename:
+        return None
+    repo = repo.strip()
+    filename = filename.strip()
+    if not repo or not filename:
+        return None
+
+    lora_dir = COMFY_DIR / "models" / "loras"
+    lora_dir.mkdir(parents=True, exist_ok=True)
+
+    try:
+        print(f"[setup] Downloading dynamic LoRA: {repo}/{filename}", flush=True)
+        local_path = hf_hub_download(
+            repo_id=repo,
+            filename=filename,
+            local_dir=str(lora_dir),
+            token=os.environ.get("HF_TOKEN"),
+        )
+        # Returns path relative to lora_dir (handles nested folders perfectly)
+        return str(Path(local_path).relative_to(lora_dir))
+    except Exception as e:
+        print(f"[error] Failed to download dynamic LoRA from {repo}/{filename}: {e}", flush=True)
+        raise RuntimeError(f"Failed to download dynamic LoRA: {e}")
+
+
+def _add_loras_to_workflow(workflow, model_node_ref, lora_name, lora_strength, lora2_name=None, lora2_strength=0.0):
+    """Inserts one or two LoraLoaderModelOnly nodes sequentially between the
+    diffusion model and the sampler."""
+    current_model = model_node_ref
+
+    # First LoRA
+    if lora_name and lora_name != LORA_NONE and float(lora_strength) != 0.0:
+        workflow["20"] = {
+            "class_type": "LoraLoaderModelOnly",
+            "inputs": {
+                "model": current_model,
+                "lora_name": lora_name,
+                "strength_model": float(lora_strength),
+            },
+        }
+        current_model = ["20", 0]
+
+    # Second dynamic LoRA
+    if lora2_name and lora2_name != LORA_NONE and float(lora2_strength) != 0.0:
+        workflow["21"] = {
+            "class_type": "LoraLoaderModelOnly",
+            "inputs": {
+                "model": current_model,
+                "lora_name": lora2_name,
+                "strength_model": float(lora2_strength),
+            },
+        }
+        current_model = ["21", 0]
+
+    return current_model
 
 
 def _build_workflow(
@@ -227,6 +267,8 @@ def _build_workflow(
     scheduler,
     lora_name=LORA_NONE,
     lora_strength=1.0,
+    lora2_name=None,
+    lora2_strength=0.0,
 ):
     workflow = {
         "1": {
@@ -285,7 +327,7 @@ def _build_workflow(
             "inputs": {"conditioning": ["2", 0]},
         }
 
-    model_ref = _add_lora_to_workflow(workflow, ["1", 0], lora_name, lora_strength)
+    model_ref = _add_loras_to_workflow(workflow, ["1", 0], lora_name, lora_strength, lora2_name, lora2_strength)
     workflow["5"]["inputs"]["model"] = model_ref
     return workflow
 
@@ -329,6 +371,8 @@ def _build_edit_workflow(
     denoise,
     lora_name=LORA_NONE,
     lora_strength=1.0,
+    lora2_name=None,
+    lora2_strength=0.0,
 ):
     workflow = {
         "1": {
@@ -391,7 +435,7 @@ def _build_edit_workflow(
             "inputs": {"conditioning": ["2", 0]},
         }
 
-    model_ref = _add_lora_to_workflow(workflow, ["1", 0], lora_name, lora_strength)
+    model_ref = _add_loras_to_workflow(workflow, ["1", 0], lora_name, lora_strength, lora2_name, lora2_strength)
     workflow["5"]["inputs"]["model"] = model_ref
     return workflow
 
@@ -424,10 +468,6 @@ def _wait_for_history(prompt_id, timeout=900):
 
 
 def _load_output_image(history_item):
-    """Downloads the generated PNG, extracts prompt/workflow from its metadata,
-    packs them into WebP's EXIF 'Make' and 'ImageDescription' tags, saves the
-    WebP file, and returns the path to it. This preserves workflow compatibility
-    with custom DnD extractors and saves tons of network bandwidth."""
     outputs = history_item.get("outputs", {})
     for output in outputs.values():
         for image in output.get("images", []):
@@ -439,16 +479,13 @@ def _load_output_image(history_item):
             response = requests.get(f"{COMFY_URL}/view", params=params, timeout=120)
             response.raise_for_status()
             
-            # 1. Save temporary PNG fetched from Comfy (with metadata inside)
             temp_png = Path("/tmp") / f"{uuid.uuid4().hex}.png"
             temp_png.write_bytes(response.content)
             
-            # 2. Extract standard ComfyUI metadata
             img = Image.open(temp_png)
             prompt_data = img.info.get("prompt", "")
             workflow_data = img.info.get("workflow", "")
             
-            # 3. Prepare EXIF payload for WebP matching your metadata ripper
             prompt_bytes = prompt_data.encode('utf-8') if isinstance(prompt_data, str) else prompt_data
             workflow_bytes = workflow_data.encode('utf-8') if isinstance(workflow_data, str) else workflow_data
             
@@ -460,14 +497,10 @@ def _load_output_image(history_item):
             }
             exif_bytes = piexif.dump(exif_dict)
             
-            # 4. Save to compressed WebP with the EXIF payload
             temp_webp = Path("/tmp") / f"{uuid.uuid4().hex}.webp"
             img.convert("RGB").save(temp_webp, "WEBP", exif=exif_bytes, quality=90)
             
-            # Cleanup temp PNG
             temp_png.unlink(missing_ok=True)
-            
-            # 5. Return filepath string so Gradio streams it byte-for-byte
             return str(temp_webp)
             
     raise RuntimeError("ComfyUI completed without returning an image.")
@@ -481,9 +514,6 @@ def _clamp_duration(value):
     return max(DURATION_MIN, min(DURATION_MAX, value))
 
 
-# spaces.GPU accepts a callable that receives the SAME positional args as the
-# decorated function and returns the requested duration. gpu_duration is just
-# the slider value passed straight through - no more 900s+megapixels guessing.
 def _duration(
     prompt,
     negative_prompt,
@@ -497,6 +527,9 @@ def _duration(
     scheduler,
     lora_name,
     lora_strength,
+    lora2_repo,
+    lora2_file,
+    lora2_strength,
     gpu_duration,
 ):
     return _clamp_duration(gpu_duration)
@@ -517,6 +550,9 @@ def _edit_duration(
     scheduler,
     lora_name,
     lora_strength,
+    edit_lora2_repo,
+    edit_lora2_file,
+    edit_lora2_strength,
     gpu_duration,
 ):
     return _clamp_duration(gpu_duration)
@@ -536,6 +572,9 @@ def generate(
     scheduler="simple",
     lora_name=LORA_NONE,
     lora_strength=1.0,
+    lora2_repo="",
+    lora2_file="",
+    lora2_strength=0.0,
     gpu_duration=DURATION_DEFAULT,
 ):
     if not prompt or not prompt.strip():
@@ -545,6 +584,10 @@ def generate(
 
     try:
         _start_comfyui()
+        
+        # Download dynamic LoRA 2 if user filled details
+        lora2_name = _download_dynamic_lora(lora2_repo, lora2_file)
+        
         workflow = _build_workflow(
             prompt,
             negative_prompt,
@@ -557,6 +600,8 @@ def generate(
             scheduler,
             lora_name,
             lora_strength,
+            lora2_name,
+            lora2_strength,
         )
         prompt_id = _queue_prompt(workflow)
         history_item = _wait_for_history(prompt_id)
@@ -581,6 +626,9 @@ def edit_image(
     scheduler="simple",
     lora_name=LORA_NONE,
     lora_strength=1.0,
+    edit_lora2_repo="",
+    edit_lora2_file="",
+    edit_lora2_strength=0.0,
     gpu_duration=DURATION_DEFAULT,
 ):
     if input_image is None:
@@ -592,6 +640,10 @@ def edit_image(
 
     try:
         _start_comfyui()
+        
+        # Download dynamic LoRA 2 if user filled details
+        lora2_name = _download_dynamic_lora(edit_lora2_repo, edit_lora2_file)
+        
         resized = _resize_for_edit(input_image, width, height)
         input_filename = _upload_image_to_comfy(resized)
         workflow = _build_edit_workflow(
@@ -606,6 +658,8 @@ def edit_image(
             denoise,
             lora_name,
             lora_strength,
+            lora2_name,
+            lora2_strength,
         )
         prompt_id = _queue_prompt(workflow)
         history_item = _wait_for_history(prompt_id)
@@ -615,8 +669,6 @@ def edit_image(
 
 
 def refresh_loras():
-    """Button handler: re-scans models/loras (via ComfyUI if it's up, plain
-    glob otherwise) and returns a Dropdown update with the fresh choices."""
     choices = _list_loras()
     return gr.update(choices=choices, value=LORA_NONE)
 
@@ -665,13 +717,20 @@ with gr.Blocks(title="Redcraft Krea2", css=CSS) as demo:
                         lora_name = gr.Dropdown(
                             choices=[LORA_NONE, LORA_FILE],
                             value=LORA_FILE,
-                            label="LoRA",
+                            label="LoRA 1",
                             info=LORA_INFO,
                             allow_custom_value=True,
                             scale=4,
                         )
                         lora_refresh = gr.Button("Refresh", scale=1)
-                    lora_strength = gr.Slider(0.0, 10.0, value=0.0, step=0.05, label="LoRA strength")
+                    lora_strength = gr.Slider(0.0, 10.0, value=0.0, step=0.05, label="LoRA 1 strength")
+                    
+                    # Spoilered dynamic LoRA 2
+                    with gr.Accordion("LoRA 2 (HuggingFace Dynamic)", open=False):
+                        lora2_repo = gr.Textbox(label="HF Repo ID", placeholder="e.g., beznogim666/test2", value="")
+                        lora2_file = gr.Textbox(label="Filename", placeholder="e.g., lora2.safetensors", value="")
+                        lora2_strength = gr.Slider(0.0, 10.0, value=0.0, step=0.05, label="LoRA 2 strength")
+
                     gpu_duration = gr.Slider(
                         DURATION_MIN,
                         DURATION_MAX,
@@ -689,7 +748,6 @@ with gr.Blocks(title="Redcraft Krea2", css=CSS) as demo:
                         scheduler = gr.Dropdown(["simple", "normal", "karras", "exponential"], value="simple", label="Scheduler")
                     run = gr.Button("Generate", variant="primary")
                 with gr.Column(scale=6):
-                    # We removed format="png" so Gradio accepts the raw WebP file path directly.
                     output = gr.Image(label="Result", elem_id="result-image")
 
             inputs = [
@@ -705,6 +763,9 @@ with gr.Blocks(title="Redcraft Krea2", css=CSS) as demo:
                 scheduler,
                 lora_name,
                 lora_strength,
+                lora2_repo,
+                lora2_file,
+                lora2_strength,
                 gpu_duration,
             ]
             run.click(generate, inputs, [output, seed])
@@ -738,13 +799,20 @@ with gr.Blocks(title="Redcraft Krea2", css=CSS) as demo:
                         edit_lora_name = gr.Dropdown(
                             choices=[LORA_NONE, LORA_FILE],
                             value=LORA_FILE,
-                            label="LoRA",
+                            label="LoRA 1",
                             info=LORA_INFO,
                             allow_custom_value=True,
                             scale=4,
                         )
                         edit_lora_refresh = gr.Button("Refresh", scale=1)
-                    edit_lora_strength = gr.Slider(-2.0, 2.0, value=1.0, step=0.05, label="LoRA strength")
+                    edit_lora_strength = gr.Slider(-2.0, 2.0, value=1.0, step=0.05, label="LoRA 1 strength")
+                    
+                    # Spoilered dynamic LoRA 2 for Edit Tab
+                    with gr.Accordion("LoRA 2 (HuggingFace Dynamic)", open=False):
+                        edit_lora2_repo = gr.Textbox(label="HF Repo ID", placeholder="e.g., beznogim666/test2", value="")
+                        edit_lora2_file = gr.Textbox(label="Filename", placeholder="e.g., lora2.safetensors", value="")
+                        edit_lora2_strength = gr.Slider(-2.0, 2.0, value=0.0, step=0.05, label="LoRA 2 strength")
+
                     edit_gpu_duration = gr.Slider(
                         DURATION_MIN,
                         DURATION_MAX,
@@ -762,7 +830,6 @@ with gr.Blocks(title="Redcraft Krea2", css=CSS) as demo:
                         edit_scheduler = gr.Dropdown(["simple", "normal", "karras", "exponential"], value="simple", label="Scheduler")
                     edit_run = gr.Button("Edit Image", variant="primary")
                 with gr.Column(scale=6):
-                    # We removed format="png" here as well.
                     edit_output = gr.Image(label="Edited image", elem_id="result-image")
 
             edit_inputs = [
@@ -780,6 +847,9 @@ with gr.Blocks(title="Redcraft Krea2", css=CSS) as demo:
                 edit_scheduler,
                 edit_lora_name,
                 edit_lora_strength,
+                edit_lora2_repo,
+                edit_lora2_file,
+                edit_lora2_strength,
                 edit_gpu_duration,
             ]
             edit_run.click(edit_image, edit_inputs, [edit_output, edit_seed])