Update app_lora1.py

app_lora1.py

@@ -2,11 +2,9 @@ import spaces
 import os
 import io
 import torch
-from PIL import Image
-
 import gradio as gr
 import requests
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, ZImagePipeline
 
 # =========================================================
 # CONFIG
@@ -21,10 +19,10 @@ MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
 os.makedirs(LOCAL_SCRIPTS_DIR, exist_ok=True)
 
 # =========================================================
-# GLOBAL STATE
+# GLOBAL STATE (CPU SAFE)
 # =========================================================
-SCRIPT_CODE = {}
-PIPELINES = {}
+SCRIPT_CODE = {}      # script_name -> code (CPU only)
+PIPELINES = {}        # script_name -> pipeline (GPU only, lazy)
 log_buffer = io.StringIO()
 
 
@@ -39,37 +37,49 @@ 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}, 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
 
+    # Quantization / GGUF
     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)
+        attn = getattr(pipe.config, "attn_implementation", None)
         if attn:
             tech.append(f"Attention: {attn}")
 
     return "\n".join(f"• {t}" for t in tech)
 
 
+# =========================================================
+# LATENT INFO
+# =========================================================
 def pipeline_debug_info(pipe):
     return f"""
 Pipeline Info
@@ -87,40 +97,7 @@ def latent_shape_info(height, width, pipe):
 
 
 # =========================================================
-# PIPELINE FEATURE REGISTRATION HELPER
-# =========================================================
-def register_pipeline_feature(pipe, text: str):
-    if not hasattr(pipe, "_enabled_features"):
-        pipe._enabled_features = []
-    pipe._enabled_features.append(text)
-
-
-# =========================================================
-# WRAPPER TO LOG ANY METHOD CALL ON PIPE OR TRANSFORMER
-# =========================================================
-def log_pipe_calls(obj, obj_name="pipe"):
-    for attr_name in dir(obj):
-        if attr_name.startswith("_"):
-            continue
-        attr = getattr(obj, attr_name)
-
-        # Skip non-callables or torch modules
-        if not callable(attr) or isinstance(attr, torch.nn.Module):
-            continue
-
-        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
+# DOWNLOAD SCRIPTS (CPU ONLY)
 # =========================================================
 def download_scripts():
     resp = requests.get(SCRIPTS_REPO_API)
@@ -143,55 +120,70 @@ SCRIPT_NAMES = download_scripts()
 
 
 # =========================================================
-# REGISTER SCRIPTS
+# REGISTER SCRIPTS (CPU ONLY)
 # =========================================================
 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 ✅"
 
 
 # =========================================================
-# BUILD PIPELINE (GPU)
+# EXTRACT LINES AFTER FROM_PRETRAINED
+# =========================================================
+def extract_pipe_lines(script_code: str):
+    lines = script_code.splitlines()
+    pipe_lines = []
+    found = False
+
+    for line in lines:
+        stripped = line.strip()
+        if not found and stripped.startswith("pipe = ZImagePipeline.from_pretrained"):
+            found = True
+            pipe_lines.append(line)
+        elif found:
+            if "pipe" in stripped:
+                pipe_lines.append(line)
+    return pipe_lines
+
+
+# =========================================================
+# GPU-ONLY PIPELINE BUILDER
 # =========================================================
 def get_pipeline(script_name):
     if script_name in PIPELINES:
         return PIPELINES[script_name]
 
     log(f"🔧 Building pipeline from {script_name}")
+    code = SCRIPT_CODE[script_name]
 
-    namespace = {
-        "__file__": script_name,
-        "__name__": "__main__",
-        "torch": torch,
-        "register_pipeline_feature": register_pipeline_feature,
-        "log_pipe_calls": log_pipe_calls,
-    }
+    # Extract lines after from_pretrained for logging
+    pipe_lines = extract_pipe_lines(code)
 
-    try:
-        exec(SCRIPT_CODE[script_name], namespace)
-    except Exception as e:
-        log(f"❌ Script failed: {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`.")
+    # Safe namespace for exec
+    namespace = {"torch": torch, "ZImagePipeline": ZImagePipeline}
 
-    pipe = namespace["pipe"]
+    pipe = None
+    for line in pipe_lines:
+        if line.strip():
+            log(f"• {line.strip()}")
+            exec(line, namespace)
+            pipe = namespace.get("pipe", 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")
+    if pipe is None:
+        raise RuntimeError(f"{script_name} did not define `pipe`.")
 
-   
+    # ZeroGPU safe
+    if not hasattr(pipe, "hf_device_map"):
+        pipe = pipe.to("cuda")
 
     PIPELINES[script_name] = pipe
-    log(f"✅ Pipeline ready: {script_name}")
-
+    log("✅ Pipeline ready")
     return pipe
 
 
@@ -216,20 +208,9 @@ def generate_image(
         raise RuntimeError("Pipeline not registered")
 
     pipe = get_pipeline(pipeline_name)
-    
-    if not hasattr(pipe, "hf_device_map"):
-      pipe = pipe.to("cuda")
 
     log("=== PIPELINE TECHNOLOGY ===")
     log(pipeline_technology_info(pipe))
-
-    log("=== PIPELINE FEATURES ===")
-    if hasattr(pipe, "_enabled_features"):
-        for f in pipe._enabled_features:
-            log(f"✔ {f}")
-    else:
-        log("✔ No explicit pipeline features registered")
-
     log("=== NEW GENERATION REQUEST ===")
     log(f"Pipeline: {pipeline_name}")
     log(f"Prompt: {prompt}")
@@ -246,6 +227,7 @@ def generate_image(
     num_images = min(max(1, int(num_images)), 3)
     generator = torch.Generator("cuda").manual_seed(int(seed))
 
+    # Run pipeline
     result = pipe(
         prompt=prompt,
         height=int(height),
@@ -258,82 +240,79 @@ def generate_image(
         output_type="pil",
     )
 
-    # 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):
-            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:
         log(pipeline_debug_info(pipe))
         log(latent_shape_info(height, width, pipe))
     except Exception as e:
         log(f"Diagnostics error: {e}")
 
-    log("Generation complete ✅")
-
-    return fixed_images, seed, log_buffer.getvalue()
+    log("✅ Generation complete")
+    return result.images, seed, log_buffer.getvalue()
 
 
 # =========================================================
-# GRADIO UI
+# GRADIO UI (original layout)
 # =========================================================
-with gr.Blocks(title="Z-Image Turbo – ZeroGPU") as demo:
-    gr.Markdown("## ⚡ Z-Image Turbo (Script-Driven · ZeroGPU Safe)")
-
-    # ------------------ Scripts selection ------------------
-    script_selector = gr.CheckboxGroup(
-        choices=SCRIPT_NAMES,
-        label="Select pipeline scripts",
-    )
-    register_btn = gr.Button("Register Scripts")
-    status = gr.Textbox(label="Status", interactive=False)
+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):
+            script_selector = gr.CheckboxGroup(
+                choices=SCRIPT_NAMES,
+                label="Select pipeline scripts"
+            )
+            register_btn = gr.Button("Register Scripts")
+            status = gr.Textbox(label="Status", interactive=False)
+
+            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)
+
+            generate_btn = gr.Button("🚀 Generate", variant="primary")
+
+        with gr.Column(scale=1):
+            pipeline_picker = gr.Dropdown(
+                choices=[],
+                label="Active Pipeline",
+            )
+            output_images = gr.Gallery(label="Generated Images", elem_classes=["gr-gallery-image"], type="pil").style(grid=[2], height="512px")
+            used_seed = gr.Number(label="Seed Used", interactive=False)
+            debug_log = gr.Textbox(
+                label="Debug Log Output",
+                lines=25,
+                interactive=False
+            )
 
     register_btn.click(
         register_scripts,
         inputs=[script_selector],
-        outputs=[status],
-    )
-
-    # ------------------ Pipeline dropdown ------------------
-    pipeline_picker = gr.Dropdown(
-        choices=[],
-        label="Active Pipeline",
+        outputs=[status]
     )
 
     register_btn.click(
         lambda s: gr.update(choices=s, value=s[0] if s else None),
         inputs=[script_selector],
-        outputs=[pipeline_picker],
+        outputs=[pipeline_picker]
     )
 
-    gr.Markdown("---")
-
-    # ------------------ Prompt + sliders ------------------
-    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")
-
-    # ------------------ Outputs ------------------
-    gallery = gr.Gallery(columns=3, height=512, object_fit="contain")
-    used_seed = gr.Number(label="Used Seed")
-    logs = gr.Textbox(lines=25, label="Logs")
-
-    run_btn.click(
+    generate_btn.click(
         generate_image,
-        inputs=[prompt, height, width, steps, seed, random_seed, images, pipeline_picker],
-        outputs=[gallery, used_seed, logs],
+        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed, num_images, pipeline_picker],
+        outputs=[output_images, used_seed, debug_log]
     )
 
 demo.queue()