Update app_lora1.py

app_lora1.py

@@ -21,10 +21,10 @@ MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
 os.makedirs(LOCAL_SCRIPTS_DIR, exist_ok=True)
 
 # =========================================================
-# GLOBAL STATE (CPU SAFE)
+# GLOBAL STATE
 # =========================================================
-SCRIPT_CODE = {}      # script_name -> code (CPU only)
-PIPELINES = {}        # script_name -> pipeline (GPU only, lazy)
+SCRIPT_CODE = {}
+PIPELINES = {}
 log_buffer = io.StringIO()
 
 
@@ -39,39 +39,29 @@ def log(msg):
 def pipeline_technology_info(pipe):
     tech = []
 
-    # Device map
     if hasattr(pipe, "hf_device_map"):
         tech.append("Device map: enabled")
     else:
         tech.append(f"Device: {pipe.device}")
 
-    # Transformer dtype
     if hasattr(pipe, "transformer"):
         try:
             tech.append(f"Transformer dtype: {pipe.transformer.dtype}")
         except Exception:
             pass
-
-        # Layerwise casting (Z-Image specific)
         if hasattr(pipe.transformer, "layerwise_casting"):
             lw = pipe.transformer.layerwise_casting
-            tech.append(
-                f"Layerwise casting: storage={lw.storage_dtype}, "
-                f"compute={lw.compute_dtype}"
-            )
+            tech.append(f"Layerwise casting: storage={lw.storage_dtype}, compute={lw.compute_dtype}")
 
-    # VAE dtype
     if hasattr(pipe, "vae"):
         try:
             tech.append(f"VAE dtype: {pipe.vae.dtype}")
         except Exception:
             pass
 
-    # GGUF / quantization
     if hasattr(pipe, "quantization_config"):
         tech.append(f"Quantization: {pipe.quantization_config}")
 
-    # Attention backend
     if hasattr(pipe, "config"):
         attn = pipe.config.get("attn_implementation", None)
         if attn:
@@ -106,7 +96,23 @@ def register_pipeline_feature(pipe, text: str):
 
 
 # =========================================================
-# DOWNLOAD SCRIPTS (CPU ONLY)
+# WRAPPER TO LOG ANY METHOD CALL ON PIPE OR TRANSFORMER
+# =========================================================
+def log_pipe_calls(obj, obj_name="pipe"):
+    for attr_name in dir(obj):
+        attr = getattr(obj, attr_name)
+        if callable(attr) and not attr_name.startswith("_"):
+            def make_wrapper(f, name):
+                def wrapper(*args, **kwargs):
+                    log(f"• {obj_name}.{name} called with args={args}, kwargs={kwargs}")
+                    return f(*args, **kwargs)
+                return wrapper
+            setattr(obj, attr_name, make_wrapper(attr, attr_name))
+    return obj
+
+
+# =========================================================
+# DOWNLOAD SCRIPTS
 # =========================================================
 def download_scripts():
     resp = requests.get(SCRIPTS_REPO_API)
@@ -129,21 +135,19 @@ SCRIPT_NAMES = download_scripts()
 
 
 # =========================================================
-# REGISTER SELECTED SCRIPTS (NO CUDA)
+# REGISTER SCRIPTS
 # =========================================================
 def register_scripts(selected_scripts):
     SCRIPT_CODE.clear()
-
     for name in selected_scripts:
         path = os.path.join(LOCAL_SCRIPTS_DIR, name)
         with open(path, "r") as f:
             SCRIPT_CODE[name] = f.read()
-
     return f"{len(SCRIPT_CODE)} script(s) registered ✅"
 
 
 # =========================================================
-# GPU-ONLY PIPELINE BUILDER (CRITICAL)
+# BUILD PIPELINE (GPU)
 # =========================================================
 def get_pipeline(script_name):
     if script_name in PIPELINES:
@@ -154,10 +158,9 @@ def get_pipeline(script_name):
     namespace = {
         "__file__": script_name,
         "__name__": "__main__",
-
-        # Minimal required globals
         "torch": torch,
         "register_pipeline_feature": register_pipeline_feature,
+        "log_pipe_calls": log_pipe_calls,
     }
 
     try:
@@ -167,19 +170,27 @@ def get_pipeline(script_name):
         raise RuntimeError(f"Pipeline build failed for {script_name}") from e
 
     if "pipe" not in namespace:
-        raise RuntimeError(
-            f"{script_name} did not define `pipe`.\n"
-            f"Each script MUST assign a variable named `pipe`."
-        )
+        raise RuntimeError(f"{script_name} did not define `pipe`.")
+
+    pipe = namespace["pipe"]
+
+    # Wrap transformer and pipe to log method calls (post-pretrained modifications)
+    if hasattr(pipe, "transformer"):
+        pipe.transformer = log_pipe_calls(pipe.transformer, "pipe.transformer")
+    pipe = log_pipe_calls(pipe, "pipe")
 
-    PIPELINES[script_name] = namespace["pipe"]
+    # ZeroGPU-safe
+    if not hasattr(pipe, "hf_device_map"):
+        pipe = pipe.to("cuda")
+
+    PIPELINES[script_name] = pipe
     log(f"✅ Pipeline ready: {script_name}")
 
-    return PIPELINES[script_name]
+    return pipe
 
 
 # =========================================================
-# IMAGE GENERATION (LOGIC UNCHANGED)
+# IMAGE GENERATION
 # =========================================================
 @spaces.GPU
 def generate_image(
@@ -200,13 +211,6 @@ def generate_image(
 
     pipe = get_pipeline(pipeline_name)
 
-    # ✅ Correct, universal, ZeroGPU-safe
-    if not hasattr(pipe, "hf_device_map"):
-        pipe = pipe.to("cuda")
-
-    # =========================================================
-    # LOG PIPELINE TECHNOLOGY AND REGISTERED FEATURES
-    # =========================================================
     log("=== PIPELINE TECHNOLOGY ===")
     log(pipeline_technology_info(pipe))
 
@@ -217,9 +221,6 @@ def generate_image(
     else:
         log("✔ No explicit pipeline features registered")
 
-    # =========================================================
-    # GENERATION LOG
-    # =========================================================
     log("=== NEW GENERATION REQUEST ===")
     log(f"Pipeline: {pipeline_name}")
     log(f"Prompt: {prompt}")
@@ -248,11 +249,15 @@ def generate_image(
         output_type="pil",
     )
 
-    # Resize images to 512x512 for Gradio
+    # Optional: scale down very large images for UI display
+    max_display_size = 1024
     fixed_images = []
     for img in result.images:
         if isinstance(img, Image.Image):
-            img = img.resize((512, 512), Image.BICUBIC)
+            w, h = img.size
+            scale = min(max_display_size / max(w, h), 1.0)
+            if scale < 1.0:
+                img = img.resize((int(w * scale), int(h * scale)), Image.BICUBIC)
         fixed_images.append(img)
 
     try:
@@ -269,68 +274,67 @@ def generate_image(
 # =========================================================
 # GRADIO UI
 # =========================================================
-with gr.Blocks(title="Z-Image Turbo – ZeroGPU") as demo:
-    gr.Markdown("## ⚡ Z-Image Turbo (Script-Driven · ZeroGPU Safe)")
+with gr.Blocks(title="Z-Image-Turbo Multi Image Demo") as demo:
+    gr.Markdown("# 🎨 Z-Image-Turbo — Multi Image")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            prompt = gr.Textbox(label="Prompt", lines=4)
+
+            with gr.Row():
+                height = gr.Slider(512, 2048, 1024, step=64, label="Height")
+                width = gr.Slider(512, 2048, 1024, step=64, label="Width")
+
+            num_images = gr.Slider(1, 3, 2, step=1, label="Number of Images")
+
+            num_inference_steps = gr.Slider(
+                1, 20, 9, step=1, label="Inference Steps",
+                info="9 steps = 8 DiT forward passes",
+            )
 
+            with gr.Row():
+                seed = gr.Number(label="Seed", value=42, precision=0)
+                randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
+
+            # Select pipeline script
+            pipeline_picker = gr.Dropdown(
+                choices=SCRIPT_NAMES,
+                value=SCRIPT_NAMES[0] if SCRIPT_NAMES else None,
+                label="Active Pipeline Script",
+            )
+
+            generate_btn = gr.Button("🚀 Generate", variant="primary")
+
+        with gr.Column(scale=1):
+            output_images = gr.Gallery(
+                label="Generated Images",
+                height=512,
+                object_fit="contain"
+            )
+            used_seed = gr.Number(label="Seed Used", interactive=False)
+            debug_log = gr.Textbox(
+                label="Debug Log Output",
+                lines=25,
+                interactive=False
+            )
+
+    register_btn = gr.Button("Register Scripts")
+    status = gr.Textbox(label="Status", interactive=False)
     script_selector = gr.CheckboxGroup(
         choices=SCRIPT_NAMES,
         label="Select pipeline scripts",
     )
 
-    register_btn = gr.Button("Register Scripts")
-    status = gr.Textbox(label="Status", interactive=False)
-
     register_btn.click(
         register_scripts,
         inputs=[script_selector],
         outputs=[status],
     )
 
-    pipeline_picker = gr.Dropdown(
-        choices=[],
-        label="Active Pipeline",
-    )
-
-    register_btn.click(
-        lambda s: gr.update(choices=s, value=s[0] if s else None),
-        inputs=[script_selector],
-        outputs=[pipeline_picker],
-    )
-
-    gr.Markdown("---")
-
-    prompt = gr.Textbox(label="Prompt", lines=3)
-    height = gr.Slider(256, 1024, 512, step=64, label="Height")
-    width = gr.Slider(256, 1024, 512, step=64, label="Width")
-    steps = gr.Slider(1, 8, 4, step=1, label="Inference Steps")
-    images = gr.Slider(1, 3, 1, step=1, label="Images")
-    seed = gr.Number(value=0, label="Seed")
-    random_seed = gr.Checkbox(value=True, label="Randomize Seed")
-
-    run_btn = gr.Button("Generate")
-
-    gallery = gr.Gallery(
-        columns=3,
-        height=512,
-        object_fit="contain",
-        label="Output (512×512)"
-    )
-    used_seed = gr.Number(label="Used Seed")
-    logs = gr.Textbox(lines=12, label="Logs")
-
-    run_btn.click(
-        generate_image,
-        inputs=[
-            prompt,
-            height,
-            width,
-            steps,
-            seed,
-            random_seed,
-            images,
-            pipeline_picker,
-        ],
-        outputs=[gallery, used_seed, logs],
+    generate_btn.click(
+        fn=generate_image,
+        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed, num_images, pipeline_picker],
+        outputs=[output_images, used_seed, debug_log],
     )
 
 demo.queue()