Fix: Load Juggernaut-Z weights via safetensors instead of broken pipeline format

RunDiffusion/Juggernaut-Z-Image only provides safetensors files, not a complete diffusers pipeline.
This fix loads the base Z-Image pipeline from Tongyi-MAI/Z-Image and then loads Juggernaut-Z
safetensors weights into the transformer component.

app.py

@@ -11,7 +11,8 @@ import gradio as gr
 from diffusers import ZImagePipeline
 
 # ==================== Configuration ====================
-MODEL_PATH = "RunDiffusion/Juggernaut-Z-Image"
+BASE_MODEL = "Tongyi-MAI/Z-Image"
+WEIGHTS_FILE = "RunDiffusion/Juggernaut-Z-Image/Juggernaut_Z_V1_by_RunDiffusion.safetensors"
 
 # ==================== Resolution Choices ====================
 RES_CHOICES = {
@@ -74,14 +75,38 @@ def get_resolution(resolution: str) -> tuple[int, int]:
 
 
 # ==================== Model Loading (Global Context) ====================
-print(f"Loading Juggernaut Z pipeline from {MODEL_PATH}...")
+print(f"Loading base Z-Image pipeline from {BASE_MODEL}...")
 pipe = ZImagePipeline.from_pretrained(
-    MODEL_PATH,
+    BASE_MODEL,
     torch_dtype=torch.bfloat16,
-    low_cpu_mem_usage=False,
 )
+
+print(f"Loading Juggernaut-Z weights from {WEIGHTS_FILE}...")
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+weights_path = hf_hub_download(
+    repo_id="RunDiffusion/Juggernaut-Z-Image",
+    filename="Juggernaut_Z_V1_by_RunDiffusion.safetensors",
+)
+state_dict = load_file(weights_path)
+
+# Load transformer weights (key component for Z-Image)
+if hasattr(pipe, 'transformer'):
+    pipe.transformer.load_state_dict(state_dict, strict=False)
+else:
+    # Alternative: load into the main unet/transformer component
+    for name, component in pipe.components.items():
+        if hasattr(component, 'load_state_dict'):
+            try:
+                component.load_state_dict(state_dict, strict=False)
+                print(f"Loaded weights into {name}")
+                break
+            except Exception as e:
+                print(f"Could not load into {name}: {e}")
+
 pipe.to("cuda")
-print("Pipeline loaded successfully!")
+print("Pipeline loaded successfully with Juggernaut-Z weights!")
 
 
 # ==================== Generation Function ====================
@@ -113,7 +138,9 @@ def generate(
     ).images[0]
     
     meta = {
-        "model": MODEL_PATH,
+        "model": "Juggernaut-Z (RunDiffusion)",
+        "base_model": BASE_MODEL,
+        "weights": "Juggernaut_Z_V1_by_RunDiffusion.safetensors",
         "prompt": prompt,
         "negative_prompt": negative_prompt,
         "resolution": f"{width} x {height}",