Update app_lora1.py

app_lora1.py

@@ -2,7 +2,6 @@ import spaces
 import os
 import io
 import sys
-import json
 import torch
 
 import gradio as gr
@@ -21,7 +20,8 @@ MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
 
 os.makedirs(LOCAL_SCRIPTS_DIR, exist_ok=True)
 
-pipe = None  # lazy-loaded pipeline
+pipelines = {}        # script_name -> pipeline
+active_pipeline = None
 log_buffer = io.StringIO()
 
 
@@ -38,7 +38,6 @@ def pipeline_debug_info(pipe):
 Pipeline Info
 -------------
 Device: {pipe.device}
-DType: {pipe.unet.dtype if hasattr(pipe, "unet") else "n/a"}
 Transformer: {pipe.transformer.__class__.__name__}
 VAE: {pipe.vae.__class__.__name__}
 """
@@ -51,7 +50,7 @@ def latent_shape_info(height, width, pipe):
 
 
 # =========================
-# SCRIPT DOWNLOADER
+# DOWNLOAD SCRIPTS
 # =========================
 def download_scripts():
     resp = requests.get(SCRIPTS_REPO_API)
@@ -74,29 +73,47 @@ SCRIPT_NAMES = download_scripts()
 
 
 # =========================
-# PIPELINE BUILDER
+# PIPELINE BUILDER (SCRIPT-DRIVEN)
 # =========================
-def build_pipeline(selected_scripts):
-    global pipe
+def build_pipelines(selected_scripts):
+    global pipelines, active_pipeline
 
-    log("Building Z-Image Turbo pipeline...")
-    log(f"Selected scripts: {selected_scripts}")
+    pipelines.clear()
+    active_pipeline = None
 
-    pipe = DiffusionPipeline.from_pretrained(
-        MODEL_ID,
-        torch_dtype=torch.bfloat16,
-        low_cpu_mem_usage=False,
-        attn_implementation="kernels-community/vllm-flash-attn3",
-    )
+    for script_name in selected_scripts:
+        log(f"\n=== Building pipeline from {script_name} ===")
+
+        script_path = os.path.join(LOCAL_SCRIPTS_DIR, script_name)
+
+        # Each script runs in its own namespace
+        namespace = {
+            "torch": torch,
+            "DiffusionPipeline": DiffusionPipeline,
+        }
+
+        with open(script_path, "r") as f:
+            code = f.read()
+
+        # Execute script (this must define `pipe`)
+        exec(code, namespace)
+
+        if "pipe" not in namespace:
+            raise RuntimeError(f"{script_name} did not create `pipe`")
 
-    pipe.to("cuda")
-    log("Pipeline loaded and moved to CUDA")
+        pipelines[script_name] = namespace["pipe"]
+        log(f"Loaded pipeline from {script_name}")
 
-    return "Pipeline ready ✅"
+    # Set default active pipeline
+    if pipelines:
+        active_pipeline = pipelines[selected_scripts[0]]
+        log(f"Active pipeline → {selected_scripts[0]}")
+
+    return f"Loaded {len(pipelines)} pipeline(s) ✅"
 
 
 # =========================
-# IMAGE GENERATION
+# IMAGE GENERATION (UNCHANGED)
 # =========================
 @spaces.GPU
 def generate_image(
@@ -107,34 +124,28 @@ def generate_image(
     seed,
     randomize_seed,
     num_images,
+    selected_pipeline_name,
 ):
-    global pipe
+    global active_pipeline
 
     log_buffer.truncate(0)
     log_buffer.seek(0)
 
-    if pipe is None:
-        raise RuntimeError("Pipeline not built yet")
+    if selected_pipeline_name not in pipelines:
+        raise RuntimeError("Pipeline not built")
+
+    pipe = pipelines[selected_pipeline_name]
 
     log("=== NEW GENERATION REQUEST ===")
+    log(f"Pipeline: {selected_pipeline_name}")
     log(f"Prompt: {prompt}")
-    log(f"Height: {height}, Width: {width}")
-    log(f"Inference Steps: {num_inference_steps}")
-    log(f"Num Images: {num_images}")
 
     if randomize_seed:
         seed = torch.randint(0, 2**32 - 1, (1,)).item()
-        log(f"Randomized Seed → {seed}")
-    else:
-        log(f"Seed: {seed}")
 
     num_images = min(max(1, int(num_images)), 3)
-
-    log(pipeline_debug_info(pipe))
-
     generator = torch.Generator("cuda").manual_seed(int(seed))
 
-    log("Running pipeline forward()...")
     result = pipe(
         prompt=prompt,
         height=int(height),
@@ -147,62 +158,61 @@ def generate_image(
         output_type="pil",
     )
 
-    try:
-        log(f"VAE latent channels: {pipe.vae.config.latent_channels}")
-        log(f"VAE scaling factor: {pipe.vae.config.scaling_factor}")
-        log(f"Transformer latent size: {pipe.transformer.config.sample_size}")
-        log(latent_shape_info(height, width, pipe))
-    except Exception as e:
-        log(f"Latent diagnostics error: {e}")
+    log(pipeline_debug_info(pipe))
+    log("Done.")
 
-    log("Pipeline finished.")
     return result.images, seed, log_buffer.getvalue()
 
 
 # =========================
 # GRADIO UI
 # =========================
-with gr.Blocks(title="Z-Image Turbo (ZeroGPU)") as demo:
-    gr.Markdown("## ⚡ Z-Image Turbo – ZeroGPU Space")
+with gr.Blocks(title="Z-Image Turbo – Script Pipelines") as demo:
+    gr.Markdown("## ⚡ Z-Image Turbo (Script-Driven Pipelines, ZeroGPU)")
 
-    with gr.Row():
-        script_selector = gr.CheckboxGroup(
-            choices=SCRIPT_NAMES,
-            label="Select Z-Image Scripts",
-        )
+    script_selector = gr.CheckboxGroup(
+        choices=SCRIPT_NAMES,
+        label="Select pipeline scripts",
+    )
 
-    build_btn = gr.Button("Build Pipeline")
-    build_status = gr.Textbox(label="Status", interactive=False)
+    build_btn = gr.Button("Build Pipelines")
+    status = gr.Textbox(label="Status")
 
     build_btn.click(
-        fn=build_pipeline,
+        build_pipelines,
         inputs=[script_selector],
-        outputs=[build_status],
+        outputs=[status],
     )
 
-    gr.Markdown("---")
+    pipeline_picker = gr.Dropdown(
+        choices=[],
+        label="Active Pipeline",
+    )
 
-    prompt = gr.Textbox(label="Prompt", lines=3)
-    with gr.Row():
-        height = gr.Slider(256, 1024, value=512, step=64, label="Height")
-        width = gr.Slider(256, 1024, value=512, step=64, label="Width")
+    build_btn.click(
+        lambda s: gr.update(choices=s, value=s[0] if s else None),
+        inputs=[script_selector],
+        outputs=[pipeline_picker],
+    )
 
-    with gr.Row():
-        steps = gr.Slider(1, 8, value=4, step=1, label="Inference Steps")
-        images = gr.Slider(1, 3, value=1, step=1, label="Images")
+    gr.Markdown("---")
 
-    with gr.Row():
-        seed = gr.Number(value=0, label="Seed")
-        random_seed = gr.Checkbox(label="Randomize Seed", value=True)
+    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")
+    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")
+    run = gr.Button("Generate")
 
-    gallery = gr.Gallery(label="Results", columns=3)
-    seed_out = gr.Number(label="Used Seed")
-    logs = gr.Textbox(label="Logs", lines=12)
+    gallery = gr.Gallery(columns=3)
+    used_seed = gr.Number(label="Used Seed")
+    logs = gr.Textbox(lines=12, label="Logs")
 
-    run_btn.click(
-        fn=generate_image,
+    run.click(
+        generate_image,
         inputs=[
             prompt,
             height,
@@ -211,8 +221,9 @@ with gr.Blocks(title="Z-Image Turbo (ZeroGPU)") as demo:
             seed,
             random_seed,
             images,
+            pipeline_picker,
         ],
-        outputs=[gallery, seed_out, logs],
+        outputs=[gallery, used_seed, logs],
     )
 
 demo.queue()