Update app_lora1.py

app_lora1.py

@@ -1,16 +1,15 @@
 import spaces
 import os
 import io
-import sys
 import torch
 
 import gradio as gr
 import requests
 from diffusers import DiffusionPipeline
 
-# =========================
+# =========================================================
 # CONFIG
-# =========================
+# =========================================================
 SCRIPTS_REPO_API = (
     "https://api.github.com/repos/asomoza/diffusers-recipes/contents/"
     "models/z-image/scripts"
@@ -20,14 +19,17 @@ MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
 
 os.makedirs(LOCAL_SCRIPTS_DIR, exist_ok=True)
 
-pipelines = {}        # script_name -> pipeline
-active_pipeline = None
+# =========================================================
+# GLOBAL STATE (CPU SAFE)
+# =========================================================
+SCRIPT_CODE = {}      # script_name -> code (CPU only)
+PIPELINES = {}        # script_name -> pipeline (GPU only, lazy)
 log_buffer = io.StringIO()
 
 
-# =========================
+# =========================================================
 # LOGGING
-# =========================
+# =========================================================
 def log(msg):
     print(msg)
     log_buffer.write(msg + "\n")
@@ -49,9 +51,9 @@ def latent_shape_info(height, width, pipe):
     return f"Expected latent size: ({h}, {w})"
 
 
-# =========================
-# DOWNLOAD SCRIPTS
-# =========================
+# =========================================================
+# DOWNLOAD SCRIPTS (CPU ONLY)
+# =========================================================
 def download_scripts():
     resp = requests.get(SCRIPTS_REPO_API)
     resp.raise_for_status()
@@ -72,49 +74,50 @@ def download_scripts():
 SCRIPT_NAMES = download_scripts()
 
 
-# =========================
-# PIPELINE BUILDER (SCRIPT-DRIVEN)
-# =========================
-def build_pipelines(selected_scripts):
-    global pipelines, active_pipeline
+# =========================================================
+# REGISTER SELECTED SCRIPTS (NO CUDA)
+# =========================================================
+def register_scripts(selected_scripts):
+    SCRIPT_CODE.clear()
 
-    pipelines.clear()
-    active_pipeline = None
+    for name in selected_scripts:
+        path = os.path.join(LOCAL_SCRIPTS_DIR, name)
+        with open(path, "r") as f:
+            SCRIPT_CODE[name] = f.read()
 
-    for script_name in selected_scripts:
-        log(f"\n=== Building pipeline from {script_name} ===")
+    return f"{len(SCRIPT_CODE)} script(s) registered ✅"
 
-        script_path = os.path.join(LOCAL_SCRIPTS_DIR, script_name)
 
-        # Each script runs in its own namespace
-        namespace = {
-            "torch": torch,
-            "DiffusionPipeline": DiffusionPipeline,
-        }
+# =========================================================
+# GPU-ONLY PIPELINE BUILDER (CRITICAL)
+# =========================================================
+@spaces.GPU
+def get_pipeline(script_name):
+    if script_name in PIPELINES:
+        return PIPELINES[script_name]
 
-        with open(script_path, "r") as f:
-            code = f.read()
+    log(f"🔧 Building pipeline from {script_name}")
 
-        # Execute script (this must define `pipe`)
-        exec(code, namespace)
+    namespace = {
+        "torch": torch,
+        "DiffusionPipeline": DiffusionPipeline,
+    }
 
-        if "pipe" not in namespace:
-            raise RuntimeError(f"{script_name} did not create `pipe`")
+    # Execute script logic (this WILL touch CUDA)
+    exec(SCRIPT_CODE[script_name], namespace)
 
-        pipelines[script_name] = namespace["pipe"]
-        log(f"Loaded pipeline from {script_name}")
+    if "pipe" not in namespace:
+        raise RuntimeError(f"{script_name} did not define `pipe`")
 
-    # Set default active pipeline
-    if pipelines:
-        active_pipeline = pipelines[selected_scripts[0]]
-        log(f"Active pipeline → {selected_scripts[0]}")
+    PIPELINES[script_name] = namespace["pipe"]
+    log(f"✅ Pipeline ready: {script_name}")
 
-    return f"Loaded {len(pipelines)} pipeline(s) ✅"
+    return PIPELINES[script_name]
 
 
-# =========================
-# IMAGE GENERATION (UNCHANGED)
-# =========================
+# =========================================================
+# IMAGE GENERATION (LOGIC UNCHANGED)
+# =========================================================
 @spaces.GPU
 def generate_image(
     prompt,
@@ -124,24 +127,28 @@ def generate_image(
     seed,
     randomize_seed,
     num_images,
-    selected_pipeline_name,
+    pipeline_name,
 ):
-    global active_pipeline
-
     log_buffer.truncate(0)
     log_buffer.seek(0)
 
-    if selected_pipeline_name not in pipelines:
-        raise RuntimeError("Pipeline not built")
+    if pipeline_name not in SCRIPT_CODE:
+        raise RuntimeError("Pipeline not registered")
 
-    pipe = pipelines[selected_pipeline_name]
+    pipe = get_pipeline(pipeline_name)
 
     log("=== NEW GENERATION REQUEST ===")
-    log(f"Pipeline: {selected_pipeline_name}")
+    log(f"Pipeline: {pipeline_name}")
     log(f"Prompt: {prompt}")
+    log(f"Height: {height}, Width: {width}")
+    log(f"Steps: {num_inference_steps}")
+    log(f"Images: {num_images}")
 
     if randomize_seed:
         seed = torch.randint(0, 2**32 - 1, (1,)).item()
+        log(f"Random Seed → {seed}")
+    else:
+        log(f"Seed → {seed}")
 
     num_images = min(max(1, int(num_images)), 3)
     generator = torch.Generator("cuda").manual_seed(int(seed))
@@ -158,28 +165,33 @@ def generate_image(
         output_type="pil",
     )
 
-    log(pipeline_debug_info(pipe))
-    log("Done.")
+    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 result.images, seed, log_buffer.getvalue()
 
 
-# =========================
+# =========================================================
 # GRADIO UI
-# =========================
-with gr.Blocks(title="Z-Image Turbo – Script Pipelines") as demo:
-    gr.Markdown("## ⚡ Z-Image Turbo (Script-Driven Pipelines, ZeroGPU)")
+# =========================================================
+with gr.Blocks(title="Z-Image Turbo – ZeroGPU") as demo:
+    gr.Markdown("## ⚡ Z-Image Turbo (Script-Driven · ZeroGPU Safe)")
 
     script_selector = gr.CheckboxGroup(
         choices=SCRIPT_NAMES,
         label="Select pipeline scripts",
     )
 
-    build_btn = gr.Button("Build Pipelines")
-    status = gr.Textbox(label="Status")
+    register_btn = gr.Button("Register Scripts")
+    status = gr.Textbox(label="Status", interactive=False)
 
-    build_btn.click(
-        build_pipelines,
+    register_btn.click(
+        register_scripts,
         inputs=[script_selector],
         outputs=[status],
     )
@@ -189,7 +201,7 @@ with gr.Blocks(title="Z-Image Turbo – Script Pipelines") as demo:
         label="Active Pipeline",
     )
 
-    build_btn.click(
+    register_btn.click(
         lambda s: gr.update(choices=s, value=s[0] if s else None),
         inputs=[script_selector],
         outputs=[pipeline_picker],
@@ -200,18 +212,18 @@ with gr.Blocks(title="Z-Image Turbo – Script Pipelines") as demo:
     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="Steps")
+    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 = gr.Button("Generate")
+    run_btn = gr.Button("Generate")
 
     gallery = gr.Gallery(columns=3)
     used_seed = gr.Number(label="Used Seed")
     logs = gr.Textbox(lines=12, label="Logs")
 
-    run.click(
+    run_btn.click(
         generate_image,
         inputs=[
             prompt,