Update app.py

app.py

@@ -27,6 +27,13 @@ COMFY_PORT = int(os.environ.get("COMFYUI_PORT", "8188"))
 COMFY_URL = f"http://{COMFY_HOST}:{COMFY_PORT}"
 MAX_SEED = 2**31 - 1
 
+# LoRA lives in the same HF repo as the main checkpoint and gets pulled down
+# at startup, same as MODEL_FILE below. "none" in the dropdown always means
+# "skip LoraLoader entirely" - everything else is downloaded automatically.
+LORA_REPO = "beznogim666/test2"
+LORA_FILE = "lora.safetensors"
+LORA_NONE = "none"
+
 # ZeroGPU per-call duration cap. Krea2-Turbo itself generates in ~10-20s once
 # ComfyUI is warm; the slider lets you dial the *requested* quota down to match
 # reality instead of the developer's original 900s+ hardcoded guess, which blew
@@ -62,7 +69,16 @@ def _validate_comfyui():
     response = requests.get(f"{COMFY_URL}/object_info", timeout=30)
     response.raise_for_status()
     object_info = response.json()
-    required_nodes = ["UNETLoader", "CLIPLoader", "VAELoader", "CLIPTextEncode", "KSampler", "VAEDecode", "SaveImage"]
+    required_nodes = [
+        "UNETLoader",
+        "CLIPLoader",
+        "VAELoader",
+        "CLIPTextEncode",
+        "KSampler",
+        "VAEDecode",
+        "SaveImage",
+        "LoraLoaderModelOnly",
+    ]
     missing = [node for node in required_nodes if node not in object_info]
     if missing:
         raise RuntimeError(f"ComfyUI is missing required nodes: {', '.join(missing)}")
@@ -77,6 +93,10 @@ def _validate_comfyui():
     if VAE_FILE not in vae_info:
         raise RuntimeError(f"Krea2 VAE is not visible to ComfyUI. First VAEs: {', '.join(vae_info[:10])}")
 
+    lora_info = object_info["LoraLoaderModelOnly"]["input"]["required"]["lora_name"][0]
+    if LORA_FILE not in lora_info:
+        raise RuntimeError(f"LoRA is not visible to ComfyUI. First LoRAs: {', '.join(lora_info[:10])}")
+
 
 def _ensure_comfyui():
     """Pure CPU/network setup: clone repo, pip install, download weights.
@@ -93,15 +113,23 @@ def _ensure_comfyui():
     diffusion_dir = COMFY_DIR / "models" / "diffusion_models"
     text_encoder_dir = COMFY_DIR / "models" / "text_encoders"
     vae_dir = COMFY_DIR / "models" / "vae"
+    lora_dir = COMFY_DIR / "models" / "loras"
     diffusion_dir.mkdir(parents=True, exist_ok=True)
     text_encoder_dir.mkdir(parents=True, exist_ok=True)
     vae_dir.mkdir(parents=True, exist_ok=True)
+    lora_dir.mkdir(parents=True, exist_ok=True)
     hf_hub_download(
         repo_id=MODEL_REPO,
         filename=MODEL_FILE,
         local_dir=str(diffusion_dir),
         token=os.environ.get("HF_TOKEN"),
     )
+    hf_hub_download(
+        repo_id=LORA_REPO,
+        filename=LORA_FILE,
+        local_dir=str(lora_dir),
+        token=os.environ.get("HF_TOKEN"),
+    )
     hf_hub_download(
         repo_id=COMFY_KREA_REPO,
         filename=f"text_encoders/{TEXT_ENCODER_FILE}",
@@ -147,7 +175,59 @@ def _start_comfyui():
         _validate_comfyui()
 
 
-def _build_workflow(prompt, negative_prompt, width, height, steps, cfg, seed, sampler, scheduler):
+def _list_loras():
+    """Query ComfyUI for whatever .safetensors currently sit in models/loras.
+    Falls back to a plain directory scan if the server isn't up yet (e.g. a
+    'Refresh' click before the first generation warmed anything up)."""
+    try:
+        response = requests.get(f"{COMFY_URL}/object_info/LoraLoaderModelOnly", timeout=5)
+        if response.ok:
+            info = response.json()
+            names = info["LoraLoaderModelOnly"]["input"]["required"]["lora_name"][0]
+            return [LORA_NONE] + list(names)
+    except Exception:
+        pass
+
+    lora_dir = COMFY_DIR / "models" / "loras"
+    if lora_dir.exists():
+        names = sorted(p.name for p in lora_dir.glob("*.safetensors"))
+        return [LORA_NONE] + names
+    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 _build_workflow(
+    prompt,
+    negative_prompt,
+    width,
+    height,
+    steps,
+    cfg,
+    seed,
+    sampler,
+    scheduler,
+    lora_name=LORA_NONE,
+    lora_strength=1.0,
+):
     workflow = {
         "1": {
             "class_type": "UNETLoader",
@@ -204,6 +284,9 @@ def _build_workflow(prompt, negative_prompt, width, height, steps, cfg, seed, sa
             "class_type": "ConditioningZeroOut",
             "inputs": {"conditioning": ["2", 0]},
         }
+
+    model_ref = _add_lora_to_workflow(workflow, ["1", 0], lora_name, lora_strength)
+    workflow["5"]["inputs"]["model"] = model_ref
     return workflow
 
 
@@ -234,7 +317,19 @@ def _resize_for_edit(image, width, height):
     return image.convert("RGB").resize((width, height), Image.LANCZOS)
 
 
-def _build_edit_workflow(input_filename, prompt, negative_prompt, steps, cfg, seed, sampler, scheduler, denoise):
+def _build_edit_workflow(
+    input_filename,
+    prompt,
+    negative_prompt,
+    steps,
+    cfg,
+    seed,
+    sampler,
+    scheduler,
+    denoise,
+    lora_name=LORA_NONE,
+    lora_strength=1.0,
+):
     workflow = {
         "1": {
             "class_type": "UNETLoader",
@@ -295,6 +390,9 @@ def _build_edit_workflow(input_filename, prompt, negative_prompt, steps, cfg, se
             "class_type": "ConditioningZeroOut",
             "inputs": {"conditioning": ["2", 0]},
         }
+
+    model_ref = _add_lora_to_workflow(workflow, ["1", 0], lora_name, lora_strength)
+    workflow["5"]["inputs"]["model"] = model_ref
     return workflow
 
 
@@ -354,7 +452,19 @@ def _clamp_duration(value):
 # 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, width, height, steps, cfg, seed, randomize_seed, sampler, scheduler, gpu_duration
+    prompt,
+    negative_prompt,
+    width,
+    height,
+    steps,
+    cfg,
+    seed,
+    randomize_seed,
+    sampler,
+    scheduler,
+    lora_name,
+    lora_strength,
+    gpu_duration,
 ):
     return _clamp_duration(gpu_duration)
 
@@ -372,6 +482,8 @@ def _edit_duration(
     randomize_seed,
     sampler,
     scheduler,
+    lora_name,
+    lora_strength,
     gpu_duration,
 ):
     return _clamp_duration(gpu_duration)
@@ -389,6 +501,8 @@ def generate(
     randomize_seed=True,
     sampler="er_sde",
     scheduler="simple",
+    lora_name=LORA_NONE,
+    lora_strength=1.0,
     gpu_duration=DURATION_DEFAULT,
 ):
     if not prompt or not prompt.strip():
@@ -398,7 +512,19 @@ def generate(
 
     try:
         _start_comfyui()
-        workflow = _build_workflow(prompt, negative_prompt, width, height, steps, cfg, seed, sampler, scheduler)
+        workflow = _build_workflow(
+            prompt,
+            negative_prompt,
+            width,
+            height,
+            steps,
+            cfg,
+            seed,
+            sampler,
+            scheduler,
+            lora_name,
+            lora_strength,
+        )
         prompt_id = _queue_prompt(workflow)
         history_item = _wait_for_history(prompt_id)
         return _load_output_image(history_item), seed
@@ -420,6 +546,8 @@ def edit_image(
     randomize_seed=True,
     sampler="er_sde",
     scheduler="simple",
+    lora_name=LORA_NONE,
+    lora_strength=1.0,
     gpu_duration=DURATION_DEFAULT,
 ):
     if input_image is None:
@@ -443,6 +571,8 @@ def edit_image(
             sampler,
             scheduler,
             denoise,
+            lora_name,
+            lora_strength,
         )
         prompt_id = _queue_prompt(workflow)
         history_item = _wait_for_history(prompt_id)
@@ -451,6 +581,12 @@ def edit_image(
         raise gr.Error(str(exc)) from exc
 
 
+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)
+
 
 CSS = """
 .gradio-container { max-width: 1120px !important; margin: 0 auto !important; }
@@ -468,6 +604,11 @@ DURATION_INFO = (
     "and the model are still loading)."
 )
 
+LORA_INFO = (
+    "Drop .safetensors files into ComfyUI/models/loras in the repo, then hit "
+    "Refresh. 'none' skips LoRA entirely."
+)
+
 with gr.Blocks(title="Redcraft Krea2", css=CSS) as demo:
     gr.Markdown("# Redcraft Krea2")
     gr.Markdown("ComfyUI-native Redcraft Krea2 generation and image editing.")
@@ -487,6 +628,17 @@ with gr.Blocks(title="Redcraft Krea2", css=CSS) as demo:
                     with gr.Row():
                         seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
                         randomize_seed = gr.Checkbox(value=True, label="Randomize seed")
+                    with gr.Row():
+                        lora_name = gr.Dropdown(
+                            choices=[LORA_NONE, LORA_FILE],
+                            value=LORA_FILE,
+                            label="LoRA",
+                            info=LORA_INFO,
+                            allow_custom_value=True,
+                            scale=4,
+                        )
+                        lora_refresh = gr.Button("Refresh", scale=1)
+                    lora_strength = gr.Slider(-2.0, 2.0, value=1.0, step=0.05, label="LoRA strength")
                     gpu_duration = gr.Slider(
                         DURATION_MIN,
                         DURATION_MAX,
@@ -517,10 +669,13 @@ with gr.Blocks(title="Redcraft Krea2", css=CSS) as demo:
                 randomize_seed,
                 sampler,
                 scheduler,
+                lora_name,
+                lora_strength,
                 gpu_duration,
             ]
             run.click(generate, inputs, [output, seed])
             prompt.submit(generate, inputs, [output, seed])
+            lora_refresh.click(refresh_loras, None, lora_name)
 
         with gr.Tab("Edit Image"):
             with gr.Row():
@@ -545,6 +700,17 @@ with gr.Blocks(title="Redcraft Krea2", css=CSS) as demo:
                     with gr.Row():
                         edit_seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
                         edit_randomize_seed = gr.Checkbox(value=True, label="Randomize seed")
+                    with gr.Row():
+                        edit_lora_name = gr.Dropdown(
+                            choices=[LORA_NONE, LORA_FILE],
+                            value=LORA_FILE,
+                            label="LoRA",
+                            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_gpu_duration = gr.Slider(
                         DURATION_MIN,
                         DURATION_MAX,
@@ -577,10 +743,13 @@ with gr.Blocks(title="Redcraft Krea2", css=CSS) as demo:
                 edit_randomize_seed,
                 edit_sampler,
                 edit_scheduler,
+                edit_lora_name,
+                edit_lora_strength,
                 edit_gpu_duration,
             ]
             edit_run.click(edit_image, edit_inputs, [edit_output, edit_seed])
             edit_prompt.submit(edit_image, edit_inputs, [edit_output, edit_seed])
+            edit_lora_refresh.click(refresh_loras, None, edit_lora_name)
 
 
 _ensure_comfyui()