Update app_quant_latent.py

app_quant_latent.py

@@ -8,6 +8,7 @@ import transformers
 import psutil
 import os
 import time
+import traceback
 
 from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
 from diffusers import ZImagePipeline, AutoModel
@@ -446,53 +447,63 @@ try:
         torch_dtype=torch_dtype,
        
     )
-    pipe.transformer.set_attention_backend("_flash_3") 
-   # pipe.load_lora_weights("bdsqlsz/qinglong_DetailedEyes_Z-Image", weight_name="qinglong_detailedeye_z-imageV2(comfy).safetensors", adapter_name="lora")
-    pipe.load_lora_weights("rahul7star/ZImageLora",
-                           weight_name="NSFW/doggystyle_pov.safetensors", adapter_name="lora")
-    
-    pipe.set_adapters(["lora",], adapter_weights=[1.])
-    pipe.fuse_lora(adapter_names=["lora"], lora_scale=0.75)
+    # If transformer supports setting backend, prefer flash-3
+    try:
+        if hasattr(pipe, "transformer") and hasattr(pipe.transformer, "set_attention_backend"):
+            pipe.transformer.set_attention_backend("_flash_3")
+            log("✅ transformer.set_attention_backend('_flash_3') called")
+    except Exception as _e:
+        log(f"⚠️ set_attention_backend failed: {_e}")
+
+    # default LoRA load (keeps your existing behaviour)
+    try:
+        pipe.load_lora_weights("rahul7star/ZImageLora",
+                               weight_name="NSFW/doggystyle_pov.safetensors", adapter_name="lora")
+        pipe.set_adapters(["lora",], adapter_weights=[1.])
+        pipe.fuse_lora(adapter_names=["lora"], lora_scale=0.75)
+    except Exception as _e:
+        log(f"⚠️ Default LoRA load failed: {_e}")
+
     debug_pipeline(pipe)
-  #  pipe.unload_lora_weights()
+    #  pipe.unload_lora_weights()
     pipe.to("cuda")
     log("✅ Pipeline built successfully.")
-    LOGS.append(log)
+    LOGS += log("Pipeline build completed.") + "\n"
 except Exception as e:
     log(f"❌ Pipeline build failed: {e}")
+    log(traceback.format_exc())
     pipe = None
 
 log_system_stats("AFTER PIPELINE BUILD")
 
 
 # -----------------------------
-# Monkey-patch prepare_latents
-# -----------------------------
+# Monkey-patch prepare_latents (safe)
 # -----------------------------
-# Monkey-patch prepare_latents
-# -----------------------------
-if pipe is not None:
+if pipe is not None and hasattr(pipe, "prepare_latents"):
     original_prepare_latents = pipe.prepare_latents
 
     def logged_prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
-        result_latents = original_prepare_latents(
-            batch_size, num_channels_latents, height, width, dtype, device, generator, latents
-        )
-        log_msg = f"🔹 prepare_latents called | shape={result_latents.shape}, dtype={result_latents.dtype}, device={result_latents.device}"
-        if hasattr(self, "_latents_log"):
-            self._latents_log.append(log_msg)
-        else:
-            self._latents_log = [log_msg]
-        return result_latents
-
-    pipe.prepare_latents = logged_prepare_latents.__get__(pipe)
-else:
-    log("❌ WARNING: Pipe not initialized; skipping prepare_latents patch")
-
-
-# Apply patch
-pipe.prepare_latents = logged_prepare_latents.__get__(pipe)
+        try:
+            result_latents = original_prepare_latents(batch_size, num_channels_latents, height, width, dtype, device, generator, latents)
+            log_msg = f"🔹 prepare_latents called | shape={result_latents.shape}, dtype={result_latents.dtype}, device={result_latents.device}"
+            if hasattr(self, "_latents_log"):
+                self._latents_log.append(log_msg)
+            else:
+                self._latents_log = [log_msg]
+            return result_latents
+        except Exception as e:
+            log(f"⚠️ prepare_latents wrapper failed: {e}")
+            raise
 
+    # apply patch safely
+    try:
+        pipe.prepare_latents = logged_prepare_latents.__get__(pipe)
+        log("✅ prepare_latents monkey-patched")
+    except Exception as e:
+        log(f"⚠️ Failed to attach prepare_latents patch: {e}")
+else:
+    log("❌ WARNING: Pipe not initialized or prepare_latents missing; skipping prepare_latents patch")
 
 
 from PIL import Image
@@ -520,7 +531,7 @@ def safe_get_latents(pipe, height, width, generator, device, LOGS):
 
 
 # --------------------------
-# Main generation function
+# Main generation function (kept exactly as your logic)
 # --------------------------
 @spaces.GPU
 def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
@@ -584,6 +595,7 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
     except Exception as e:
         LOGS.append(f"❌ Total failure: {e}")
         return placeholder, [placeholder], LOGS
+
 @spaces.GPU
 def generate_image_backup(prompt, height, width, steps, seed, guidance_scale=0.0, return_latents=False):
     """
@@ -676,35 +688,35 @@ def generate_image_backup(prompt, height, width, steps, seed, guidance_scale=0.0
 # UI
 # ============================================================
 
-# with gr.Blocks(title="Z-Image- experiment - dont run")as demo:
-#   gr.Markdown("# **🚀 do not run Z-Image-Turbo — Final Image & Latents**")
-
-
-#   with gr.Row():
-#     with gr.Column(scale=1):
-#         prompt = gr.Textbox(label="Prompt", value="boat in Ocean")
-#         height = gr.Slider(256, 2048, value=1024, step=8, label="Height")
-#         width = gr.Slider(256, 2048, value=1024, step=8, label="Width")
-#         steps = gr.Slider(1, 50, value=20, step=1, label="Inference Steps")
-#         seed = gr.Number(value=42, label="Seed")
-#         run_btn = gr.Button("Generate Image")
-
-#     with gr.Column(scale=1):
-#         final_image = gr.Image(label="Final Image")
-#         latent_gallery = gr.Gallery(
-#            label="Latent Steps",
-#                 columns=4,
-#               height=256,
-#              preview=True
-#               )
-
-#         logs_box = gr.Textbox(label="Logs", lines=15)
+# Utility: scan local HF cache for safetensors in a repo folder name
+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.
+    """
+    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:
+    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("/", "_")
+    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
+    safe_list = sorted(list(dict.fromkeys(safe_list)))
+    return safe_list
 
-#     run_btn.click(
-#       generate_image,
-#       inputs=[prompt, height, width, steps, seed],
-#       outputs=[final_image, latent_gallery, logs_box]
-#      )
 
 with gr.Blocks(title="Z-Image-Turbo") as demo:
     with gr.Tabs():
@@ -727,6 +739,84 @@ with gr.Blocks(title="Z-Image-Turbo") as demo:
         with gr.TabItem("Logs"):
             logs_box = gr.Textbox(label="All Logs", lines=25)
 
+    # New UI: LoRA repo textbox, dropdown, refresh & rebuild
+    with gr.Row():
+        lora_repo = gr.Textbox(label="LoRA Repo (HF id)", value="rahul7star/ZImageLora", placeholder="e.g. rahul7star/ZImageLora")
+        lora_dropdown = gr.Dropdown(choices=[], label="LoRA files (from local cache)")
+        refresh_lora_btn = gr.Button("Refresh LoRA List")
+        rebuild_pipe_btn = gr.Button("Rebuild pipeline (use selected LoRA)")
+
+    # Refresh callback: repopulate dropdown from repo text
+    def refresh_lora_list(repo_name):
+        try:
+            files = list_loras_from_repo(repo_name)
+            if not files:
+                return gr.update(choices=[], value=None)
+            return gr.update(choices=files, value=files[0])
+        except Exception as e:
+            log(f"⚠️ refresh_lora_list failed: {e}")
+            return gr.update(choices=[], value=None)
+
+    refresh_lora_btn.click(refresh_lora_list, inputs=[lora_repo], outputs=[lora_dropdown])
+
+    # Rebuild callback: build pipeline with selected lora file path (if any)
+    def rebuild_pipeline_with_lora(lora_path, repo_name):
+        global pipe, LOGS
+        try:
+            log(f"🔄 Rebuilding pipeline using LoRA repo={repo_name} file={lora_path}")
+            # call existing logic to rebuild: attempt to create new pipeline then load lora file
+            pipe = ZImagePipeline.from_pretrained(
+                model_id,
+                transformer=transformer,
+                text_encoder=text_encoder,
+                torch_dtype=torch_dtype,
+            )
+            # try set backend
+            try:
+                if hasattr(pipe, "transformer") and hasattr(pipe.transformer, "set_attention_backend"):
+                    pipe.transformer.set_attention_backend("_flash_3")
+            except Exception as _e:
+                log(f"⚠️ set_attention_backend failed during rebuild: {_e}")
+
+            # load selected lora if provided
+            if lora_path:
+                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),
+                                           adapter_name="lora")
+                    pipe.set_adapters(["lora"], adapter_weights=[1.])
+                    pipe.fuse_lora(adapter_names=["lora"], lora_scale=0.75)
+                except Exception as _e:
+                    log(f"⚠️ Failed to load selected LoRA during rebuild: {_e}")
+
+            # finalize
+            debug_pipeline(pipe)
+            pipe.to("cuda")
+            # re-attach monkey patch safely
+            if pipe is not None and hasattr(pipe, "prepare_latents"):
+                try:
+                    original_prepare = pipe.prepare_latents
+                    def logged_prepare(self, *args, **kwargs):
+                        lat = original_prepare(*args, **kwargs)
+                        msg = f"🔹 prepare_latents called | shape={lat.shape}, dtype={lat.dtype}"
+                        if hasattr(self, "_latents_log"):
+                            self._latents_log.append(msg)
+                        else:
+                            self._latents_log = [msg]
+                        return lat
+                    pipe.prepare_latents = logged_prepare.__get__(pipe)
+                    log("✅ Re-applied prepare_latents monkey patch after rebuild")
+                except Exception as _e:
+                    log(f"⚠️ Could not re-apply prepare_latents patch: {_e}")
+            return "\n".join([LOGS, "Rebuild complete."])
+        except Exception as e:
+            log(f"❌ Rebuild pipeline failed: {e}")
+            log(traceback.format_exc())
+            return "\n".join([LOGS, f"Rebuild failed: {e}"])
+
+    rebuild_pipe_btn.click(rebuild_pipeline_with_lora, inputs=[lora_dropdown, lora_repo], outputs=[logs_box])
+
     # Wire the button AFTER all components exist
     run_btn.click(
         generate_image,
@@ -734,9 +824,4 @@ with gr.Blocks(title="Z-Image-Turbo") as demo:
         outputs=[final_image, latent_gallery, logs_box]
     )
 
-
-
-
-
-
-demo.launch()
+demo.launch()