Update app_quant_latent.py

app_quant_latent.py

@@ -693,27 +693,70 @@ def list_loras_from_repo(repo_id):
     """
     Attempts to find safetensors inside HF cache directory for repo_id.
     This only scans local cache; it does NOT download anything.
+
+    Returns:
+      A list of strings suitable for showing in the dropdown. Prefer returning
+      paths relative to the repo root (e.g. "NSFW/doggystyle_pov.safetensors") so that
+      pipe.load_lora_weights(repo_id, weight_name=that_path) works for nested files.
+      If a relative path can't be determined, returns absolute cached file paths.
     """
     if not repo_id:
         return []
-    # Map a repo id to local cache folder name heuristic (works for many cases)
+
     safe_list = []
-    # Common Hugging Face cache root:
+
+    # Candidate cache roots
     hf_cache = os.path.expanduser("~/.cache/huggingface/hub")
-    # Also check /home/user/.cache/huggingface/hub (Spaces environments)
     alt_cache = "/home/user/.cache/huggingface/hub"
-
     candidates = [hf_cache, alt_cache]
-    needle = repo_id.replace("/", "_")
+
+    # Normalize repo variants to search for in path
+    owner_repo = repo_id.replace("/", "_")
+    owner_repo_dash = repo_id.replace("/", "-")
+    owner_repo_double = repo_id.replace("/", "--")
+
+    # Walk caches and collect safetensors
     for root_cache in candidates:
         if not os.path.exists(root_cache):
             continue
-        for root, dirs, files in os.walk(root_cache):
-            if needle in root:
-                for f in files:
-                    if f.endswith(".safetensors"):
-                        safe_list.append(os.path.join(root, f))
-    # de-duplicate and sort
+        for dirpath, dirnames, filenames in os.walk(root_cache):
+            for f in filenames:
+                if not f.endswith(".safetensors"):
+                    continue
+                full_path = os.path.join(dirpath, f)
+
+                # try to find a repo-root-like substring in dirpath
+                chosen_base = None
+                for pattern in (owner_repo_double, owner_repo_dash, owner_repo):
+                    idx = dirpath.find(pattern)
+                    if idx != -1:
+                        chosen_base = dirpath[: idx + len(pattern)]
+                        break
+
+                # fallback: look for the repo folder name (last component) e.g., "ZImageLora"
+                if chosen_base is None:
+                    repo_tail = repo_id.split("/")[-1]
+                    idx2 = dirpath.find(repo_tail)
+                    if idx2 != -1:
+                        chosen_base = dirpath[: idx2 + len(repo_tail)]
+
+                # If we found a base that looks like the cached repo root, compute relative path
+                if chosen_base:
+                    try:
+                        rel = os.path.relpath(full_path, chosen_base)
+                        # If relpath goes up (starts with ..) then prefer full_path
+                        if rel and not rel.startswith(".."):
+                            # Normalize to forward slashes for HF repo weight_name usage
+                            rel_normalized = rel.replace(os.sep, "/")
+                            safe_list.append(rel_normalized)
+                            continue
+                    except Exception:
+                        pass
+
+                # Otherwise append absolute path (last resort)
+                safe_list.append(full_path)
+
+    # remove duplicates and sort
     safe_list = sorted(list(dict.fromkeys(safe_list)))
     return safe_list
 
@@ -780,15 +823,62 @@ with gr.Blocks(title="Z-Image-Turbo") as demo:
 
             # load selected lora if provided
             if lora_path:
+                weight_name_to_use = None
+
+                # If dropdown provided a relative-style path (contains a slash or no leading /),
+                # use it directly as weight_name (HF expects "path/inside/repo.safetensors")
+                if ("/" in lora_path) and not os.path.isabs(lora_path):
+                    weight_name_to_use = lora_path
+                else:
+                    # It might be an absolute path in cache; try to compute relative path to repo cache root
+                    abs_path = lora_path if os.path.isabs(lora_path) else None
+                    if abs_path and os.path.exists(abs_path):
+                        # attempt to find repo-root-ish substring in abs_path
+                        repo_variants = [
+                            repo_name.replace("/", "--"),
+                            repo_name.replace("/", "-"),
+                            repo_name.replace("/", "_"),
+                            repo_name.split("/")[-1],
+                        ]
+                        chosen_base = None
+                        for v in repo_variants:
+                            idx = abs_path.find(v)
+                            if idx != -1:
+                                chosen_base = abs_path[: idx + len(v)]
+                                break
+                        if chosen_base:
+                            try:
+                                rel = os.path.relpath(abs_path, chosen_base)
+                                if rel and not rel.startswith(".."):
+                                    weight_name_to_use = rel.replace(os.sep, "/")
+                            except Exception:
+                                weight_name_to_use = None
+
+                    # fallback to basename
+                    if weight_name_to_use is None:
+                        weight_name_to_use = os.path.basename(lora_path)
+
+                # Now attempt to load
                 try:
-                    # repo_name must be HF repo id where load_lora_weights expects it; if user provided repo id use that
                     pipe.load_lora_weights(repo_name or "rahul7star/ZImageLora",
-                                           weight_name=os.path.basename(lora_path),
+                                           weight_name=weight_name_to_use,
                                            adapter_name="lora")
                     pipe.set_adapters(["lora"], adapter_weights=[1.])
                     pipe.fuse_lora(adapter_names=["lora"], lora_scale=0.75)
+                    log(f"✅ Loaded LoRA weight: {weight_name_to_use} from repo {repo_name}")
                 except Exception as _e:
-                    log(f"⚠️ Failed to load selected LoRA during rebuild: {_e}")
+                    log(f"⚠️ Failed to load selected LoRA during rebuild using weight_name='{weight_name_to_use}': {_e}")
+                    # as last resort, try loading using basename
+                    try:
+                        fallback_name = os.path.basename(lora_path)
+                        pipe.load_lora_weights(repo_name or "rahul7star/ZImageLora",
+                                               weight_name=fallback_name,
+                                               adapter_name="lora")
+                        pipe.set_adapters(["lora"], adapter_weights=[1.])
+                        pipe.fuse_lora(adapter_names=["lora"], lora_scale=0.75)
+                        log(f"✅ Fallback loaded LoRA weight basename: {fallback_name}")
+                    except Exception as _e2:
+                        log(f"❌ Fallback LoRA load also failed: {_e2}")
 
             # finalize
             debug_pipeline(pipe)