add initial lora loading

app.py

@@ -19,6 +19,7 @@ import torch
 from huggingface_hub import hf_hub_download
 
 from diffusers import Ideogram4Pipeline
+from lora_utils import ensure_loras_loaded
 
 # Runtime shim (keeps the bundled diffusers pristine): cu130-era bitsandbytes returns Params4bit.shape as a
 # plain tuple, but diffusers' check_quantized_param_shape calls .numel() on it. math.prod handles both, so
@@ -43,6 +44,7 @@ LM_HEAD_REPO = "multimodalart/qwen3-vl-8b-instruct-lm-head"
 AOTI_REPO = "multimodalart/i4-block-aoti"
 AOTI_BLOCK_FILE = "Ideogram4TransformerBlock/package.pt2"
 MAX_SEED = 2**31 - 1
+LAST_LORAS = {}
 
 # Prompt upsampling: Ideogram's hosted magic-prompt (default) with the local Qwen graft as fallback.
 IDEOGRAM_MAGIC_PROMPT_URL = "https://api.ideogram.ai/v1/ideogram-v4/magic-prompt"
@@ -143,16 +145,48 @@ def _per_step(width, height):
     return max(0.2, _PS_A + _PS_B * ((int(width) // 16) * (int(height) // 16)))
 
 
-def _gpu_duration(final_prompt, mode, width, height, seed, do_local, progress=None):
+def _gpu_duration(
+    final_prompt,
+    mode,
+    width,
+    height,
+    seed,
+    do_local,
+    is_raw_json=False,
+    lora_spec="",
+    lora_scale=1.0,
+    progress=None,
+):
+    lora_requested = bool(lora_spec and str(lora_spec).strip())
+    if is_raw_json:
+        return 80 if lora_requested else 20
     steps = MODES.get(mode, MODES["Default · 20 steps"])["num_inference_steps"]
     budget = steps * _per_step(width, height) + DIFFUSION_OVERHEAD_S
     if do_local:
         budget += LOCAL_UPSAMPLE_S
-    return max(60, int(math.ceil(budget * DURATION_MARGIN)))
+    if lora_requested:
+        budget += 10
+    return max(20, int(math.ceil(budget * DURATION_MARGIN)))
 
 
 @spaces.GPU(duration=_gpu_duration, size="xlarge")
-def _gpu_generate(final_prompt, mode, width, height, seed, do_local, progress=gr.Progress(track_tqdm=True)):
+def _gpu_generate(
+    final_prompt,
+    mode,
+    width,
+    height,
+    seed,
+    do_local,
+    is_raw_json=False,
+    lora_spec="",
+    lora_scale=1.0,
+    progress=gr.Progress(track_tqdm=True),
+):
+    try:
+        active_loras = ensure_loras_loaded(pipe, lora_spec, float(lora_scale), LAST_LORAS)
+    except Exception as e:
+        raise gr.Error(f"Failed to load LoRA adapters: {e}") from e
+
     # Overlap the AOTI block-patch with the (transformer-idle) local upsample, if any.
     aoti_thread = Thread(target=_apply_aoti, daemon=True)
     aoti_thread.start()
@@ -180,6 +214,15 @@ def _gpu_generate(final_prompt, mode, width, height, seed, do_local, progress=gr
         caption = json.loads(final_prompt)
     except Exception:
         caption = {"prompt": final_prompt}
+    if active_loras:
+        caption["loras"] = [
+            {
+                "repo": entry["repo_id"],
+                "weight_name": entry["weight_name"],
+                "scale": entry["scale"],
+            }
+            for entry in active_loras
+        ]
     return image, int(seed), caption
 
 
@@ -192,8 +235,12 @@ def generate(
     seed=0,
     randomize_seed=False,
     raw_json_prompt="",
+    lora_spec="",
+    lora_scale=1.0,
     progress=gr.Progress(track_tqdm=True),
 ):
+    if isinstance(lora_scale, str):
+        lora_scale = float(lora_scale)
     if randomize_seed or seed < 0:
         seed = random.randint(0, MAX_SEED)
 
@@ -203,7 +250,7 @@ def generate(
             final_prompt = json.dumps(json.loads(raw_json_prompt), ensure_ascii=False, separators=(",", ":"))
         except Exception as e:
             raise gr.Error(f"Raw JSON prompt is not valid JSON: {e}")
-        return _gpu_generate(final_prompt, mode, width, height, seed, False)
+        return _gpu_generate(final_prompt, mode, width, height, seed, False, True, lora_spec, lora_scale)
 
     # Remote upsample is a network call -> run it here, OFF the GPU. Fall back to local (on-GPU) on failure.
     final_prompt, do_local = prompt, True
@@ -218,7 +265,7 @@ def generate(
             print(f"[upsample] remote failed, falling back to local: {e!r}", flush=True)
             gr.Warning("Ideogram API unavailable — using the local Qwen upsampler.")
 
-    return _gpu_generate(final_prompt, mode, width, height, seed, do_local)
+    return _gpu_generate(final_prompt, mode, width, height, seed, do_local, False, lora_spec, lora_scale)
 
 
 @spaces.GPU(size="xlarge")
@@ -267,6 +314,19 @@ with gr.Blocks(theme=gr.themes.Citrus(), title="Ideogram 4", css=CSS) as demo:
                     placeholder='{"prompt":"a ginger cat wearing a tiny wizard hat reading a spellbook"}',
                     buttons=["copy"],
                 )
+                lora_spec = gr.Textbox(
+                    label="LoRAs (one per line)",
+                    info="Format: username/repo:weights.safetensors@0.8 or a direct Hugging Face .safetensors URL. Leave empty for no LoRA.",
+                    lines=4,
+                    placeholder="vladi/real-loras:ema4_flux2_klein_9b_000002000.safetensors\nvladi/loras:klein9b/klein_snofs_v1_4.safetensors@0.8",
+                )
+                lora_scale = gr.Slider(
+                    0.0,
+                    2.0,
+                    value=1.0,
+                    step=0.05,
+                    label="LoRA scale multiplier",
+                )
                 with gr.Row():
                     width = gr.Slider(512, 2048, value=1024, step=64, label="Width")
                     height = gr.Slider(512, 2048, value=1024, step=64, label="Height")
@@ -292,7 +352,7 @@ with gr.Blocks(theme=gr.themes.Citrus(), title="Ideogram 4", css=CSS) as demo:
 
     run.click(
         generate,
-        inputs=[prompt, mode, upsampler, width, height, seed, randomize, raw_json_prompt],
+        inputs=[prompt, mode, upsampler, width, height, seed, randomize, raw_json_prompt, lora_spec, lora_scale],
         outputs=[out_image, seed, out_caption],
     )
 

lora_utils.py

@@ -0,0 +1,270 @@
+import hashlib
+import os
+from urllib.parse import urlparse
+
+
+ADAPTER_NAME_PREFIX = "custom"
+HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("hf")
+
+
+def _parse_hf_lora_url(url: str):
+    parsed = urlparse(url)
+    if "huggingface.co" not in parsed.netloc:
+        return None, None
+
+    path_parts = [part for part in parsed.path.split("/") if part]
+    if len(path_parts) < 2:
+        return None, None
+
+    repo_id = f"{path_parts[0]}/{path_parts[1]}"
+    weight_parts = path_parts[2:]
+    if len(weight_parts) >= 2 and weight_parts[0] in {"blob", "resolve"}:
+        weight_parts = weight_parts[2:]
+    weight_name = "/".join(weight_parts) if weight_parts else None
+    if not weight_name or not weight_name.endswith(".safetensors"):
+        return repo_id, None
+    return repo_id, weight_name
+
+
+def _split_lora_spec(spec: str):
+    if not spec:
+        return None, None
+
+    spec = spec.strip()
+    if not spec:
+        return None, None
+
+    if spec.startswith("http://") or spec.startswith("https://"):
+        return _parse_hf_lora_url(spec)
+    if ":" in spec:
+        repo_id, weight_name = spec.split(":", 1)
+        return repo_id.strip(), weight_name.strip()
+    return spec, None
+
+
+def _split_adapter_line_scale(line: str):
+    if "@" not in line:
+        return line, 1.0
+
+    spec_candidate, scale_candidate = line.rsplit("@", 1)
+    try:
+        inline_scale = float(scale_candidate.strip())
+    except ValueError:
+        return line, 1.0
+    return spec_candidate.strip(), inline_scale
+
+
+def parse_adapter_specs(spec_text: str, global_scale: float):
+    if not spec_text or not spec_text.strip():
+        return []
+
+    requested_entries = []
+    seen_keys = set()
+
+    for line_number, raw_line in enumerate(spec_text.splitlines(), start=1):
+        line = raw_line.strip()
+        if not line:
+            continue
+
+        spec, inline_scale = _split_adapter_line_scale(line)
+        repo_id, weight_name = _split_lora_spec(spec)
+        if not repo_id or not weight_name:
+            raise ValueError(
+                "Please provide LoRA entries as "
+                "'user/repo:weights.safetensors' or direct .safetensors URLs. "
+                f"Invalid line {line_number}: {raw_line!r}"
+            )
+
+        adapter_key = (repo_id, weight_name)
+        if adapter_key in seen_keys:
+            raise ValueError(
+                f"Duplicate LoRA entry for '{repo_id}:{weight_name}' on line {line_number}."
+            )
+        seen_keys.add(adapter_key)
+
+        requested_entries.append(
+            {
+                "key": adapter_key,
+                "repo_id": repo_id,
+                "weight_name": weight_name,
+                "adapter_name": adapter_runtime_name(adapter_key),
+                "inline_scale": inline_scale,
+                "global_scale": global_scale,
+                "scale": inline_scale * global_scale,
+            }
+        )
+
+    return requested_entries
+
+
+def adapter_runtime_name(adapter_key):
+    digest = hashlib.sha1(f"{adapter_key[0]}:{adapter_key[1]}".encode("utf-8")).hexdigest()[:12]
+    return f"{ADAPTER_NAME_PREFIX}_{digest}"
+
+
+def _iter_adapter_hosts(pipe):
+    seen = set()
+    for host in (
+        pipe,
+        getattr(pipe, "transformer", None),
+        getattr(pipe, "unconditional_transformer", None),
+    ):
+        if host is None or id(host) in seen:
+            continue
+        seen.add(id(host))
+        yield host
+
+
+def _flatten_adapter_names(adapter_mapping):
+    if isinstance(adapter_mapping, dict):
+        names = set()
+        for adapters in adapter_mapping.values():
+            if isinstance(adapters, str):
+                names.add(adapters)
+            else:
+                names.update(adapters)
+        return names
+    if isinstance(adapter_mapping, str):
+        return {adapter_mapping}
+    if adapter_mapping is None:
+        return set()
+    return set(adapter_mapping)
+
+
+def _sorted_lora_entries(entries):
+    return sorted(entries, key=lambda entry: entry["adapter_name"])
+
+
+def safe_unload_lora_adapters(pipe):
+    deleted = False
+    for host in _iter_adapter_hosts(pipe):
+        if not hasattr(host, "delete_adapters"):
+            continue
+        try:
+            adapter_names = sorted(_flatten_adapter_names(host.get_list_adapters()))
+        except Exception:
+            adapter_names = []
+        for adapter_name in adapter_names:
+            try:
+                host.delete_adapters(adapter_name)
+                deleted = True
+            except Exception:
+                pass
+    if deleted:
+        return
+
+    if hasattr(pipe, "unload_lora_weights"):
+        try:
+            pipe.unload_lora_weights()
+            return
+        except Exception:
+            pass
+
+    for host in _iter_adapter_hosts(pipe):
+        if hasattr(host, "set_adapters"):
+            try:
+                host.set_adapters([])
+            except Exception:
+                pass
+        if hasattr(host, "disable_adapters"):
+            try:
+                host.disable_adapters()
+            except Exception:
+                pass
+        if hasattr(host, "disable_lora"):
+            try:
+                host.disable_lora()
+            except Exception:
+                pass
+
+
+def _set_adapters_on_host(host, adapter_names, adapter_weights):
+    if not hasattr(host, "set_adapters"):
+        return False
+
+    if not adapter_names:
+        try:
+            host.set_adapters([])
+            return True
+        except Exception:
+            return False
+
+    for kwargs in (
+        {"adapter_weights": adapter_weights},
+        {"weights": adapter_weights},
+    ):
+        try:
+            host.set_adapters(adapter_names, **kwargs)
+            return True
+        except TypeError:
+            continue
+        except Exception:
+            return False
+    return False
+
+
+def apply_lora_adapters(pipe, lora_entries):
+    if not lora_entries:
+        safe_unload_lora_adapters(pipe)
+        return
+
+    sorted_entries = _sorted_lora_entries(lora_entries)
+    adapter_names = [entry["adapter_name"] for entry in sorted_entries]
+    adapter_weights = [entry["scale"] for entry in sorted_entries]
+
+    for host in _iter_adapter_hosts(pipe):
+        if _set_adapters_on_host(host, adapter_names, adapter_weights):
+            return
+
+    if len(adapter_names) == 1 and hasattr(pipe, "set_lora_scale"):
+        pipe.set_lora_scale(adapter_weights[0])
+        return
+
+    raise ValueError("This runtime does not support activating multiple LoRA adapters.")
+
+
+def load_lora_adapter(pipe, entry, token=HF_TOKEN):
+    load_kwargs = {
+        "weight_name": entry["weight_name"],
+        "adapter_name": entry["adapter_name"],
+    }
+    if token:
+        load_kwargs["token"] = token
+
+    try:
+        pipe.load_lora_weights(entry["repo_id"], **load_kwargs)
+    except TypeError:
+        if not token:
+            raise
+        load_kwargs.pop("token", None)
+        pipe.load_lora_weights(entry["repo_id"], **load_kwargs)
+
+
+def ensure_loras_loaded(pipe, spec_text: str, global_scale: float, active_by_key: dict, token=HF_TOKEN):
+    desired_entries = parse_adapter_specs(spec_text, global_scale)
+    desired_by_key = {entry["key"]: entry for entry in desired_entries}
+
+    if not desired_entries:
+        if active_by_key:
+            safe_unload_lora_adapters(pipe)
+            active_by_key.clear()
+        return []
+
+    if set(active_by_key.keys()) != set(desired_by_key.keys()):
+        try:
+            safe_unload_lora_adapters(pipe)
+            loaded_entries = []
+            for entry in _sorted_lora_entries(desired_entries):
+                load_lora_adapter(pipe, entry, token=token)
+                loaded_entries.append(entry)
+            apply_lora_adapters(pipe, loaded_entries)
+        except Exception:
+            safe_unload_lora_adapters(pipe)
+            active_by_key.clear()
+            raise
+    else:
+        apply_lora_adapters(pipe, desired_entries)
+
+    active_by_key.clear()
+    active_by_key.update(desired_by_key)
+    return desired_entries

tests/test_lora_utils.py

@@ -0,0 +1,53 @@
+import unittest
+
+from lora_utils import parse_adapter_specs
+
+
+class ParseAdapterSpecsTest(unittest.TestCase):
+    def test_blank_spec_returns_no_entries(self):
+        self.assertEqual(parse_adapter_specs("\n  \n", 1.0), [])
+
+    def test_parses_repo_weight_and_default_scale(self):
+        entries = parse_adapter_specs("vladi/real-loras:ema4_flux2_klein_9b_000002000.safetensors", 1.0)
+
+        self.assertEqual(entries[0]["repo_id"], "vladi/real-loras")
+        self.assertEqual(entries[0]["weight_name"], "ema4_flux2_klein_9b_000002000.safetensors")
+        self.assertEqual(entries[0]["scale"], 1.0)
+
+    def test_parses_multiple_entries_and_inline_scale(self):
+        entries = parse_adapter_specs(
+            "vladi/real-loras:ema4_flux2_klein_9b_000002000.safetensors\n"
+            "vladi/loras:klein9b/klein_snofs_v1_4.safetensors@0.8",
+            0.5,
+        )
+
+        self.assertEqual(len(entries), 2)
+        self.assertEqual(entries[1]["repo_id"], "vladi/loras")
+        self.assertEqual(entries[1]["weight_name"], "klein9b/klein_snofs_v1_4.safetensors")
+        self.assertEqual(entries[1]["scale"], 0.4)
+
+    def test_parses_huggingface_url_with_subfolder(self):
+        entries = parse_adapter_specs(
+            "https://huggingface.co/vladi/loras/blob/dev/klein9b/klein_snofs_v1_4.safetensors@0.8",
+            1.0,
+        )
+
+        self.assertEqual(entries[0]["repo_id"], "vladi/loras")
+        self.assertEqual(entries[0]["weight_name"], "klein9b/klein_snofs_v1_4.safetensors")
+        self.assertEqual(entries[0]["scale"], 0.8)
+
+    def test_rejects_missing_weight_name(self):
+        with self.assertRaises(ValueError):
+            parse_adapter_specs("vladi/real-loras", 1.0)
+
+    def test_rejects_duplicate_entries(self):
+        with self.assertRaises(ValueError):
+            parse_adapter_specs(
+                "vladi/real-loras:ema4_flux2_klein_9b_000002000.safetensors\n"
+                "vladi/real-loras:ema4_flux2_klein_9b_000002000.safetensors@0.8",
+                1.0,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()