Update app_lora1.py

app_lora1.py

@@ -1,316 +1,218 @@
+import os
+import io
+import sys
+import json
+import torch
 import spaces
 import gradio as gr
-import torch
-import os
-import traceback
-from diffusers import ZImagePipeline
-from huggingface_hub import list_repo_files
-from PIL import Image
-
+import requests
+from diffusers import DiffusionPipeline
 
-# ============================================================
+# =========================
 # CONFIG
-# ============================================================
-
+# =========================
+SCRIPTS_REPO_API = (
+    "https://api.github.com/repos/asomoza/diffusers-recipes/contents/"
+    "models/z-image/scripts"
+)
+LOCAL_SCRIPTS_DIR = "z_image_scripts"
 MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
-DEFAULT_LORA_REPO = "rahul7star/ZImageLora"
 
-DTYPE = torch.bfloat16
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+os.makedirs(LOCAL_SCRIPTS_DIR, exist_ok=True)
 
-# ============================================================
-# GLOBAL STATE
-# ============================================================
+pipe = None  # lazy-loaded pipeline
+log_buffer = io.StringIO()
 
-pipe = None
-CURRENT_LORA_REPO = None
-CURRENT_LORA_FILE = None
 
-# ============================================================
+# =========================
 # LOGGING
-# ============================================================
-
+# =========================
 def log(msg):
     print(msg)
-    return msg
+    log_buffer.write(msg + "\n")
 
-# ============================================================
-# PIPELINE BUILD (ONCE)
-# ============================================================
 
-try:
-    pipe = ZImagePipeline.from_pretrained(
-        MODEL_ID,
-        torch_dtype=DTYPE,
-    )
+def pipeline_debug_info(pipe):
+    return f"""
+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__}
+"""
 
-    
 
-    pipe.to(DEVICE)
-    log("✅ Pipeline built successfully")
+def latent_shape_info(height, width, pipe):
+    h = height // pipe.vae_scale_factor
+    w = width // pipe.vae_scale_factor
+    return f"Expected latent size: ({h}, {w})"
 
-except Exception as e:
-    log("❌ Pipeline build failed")
-    log(traceback.format_exc())
-    pipe = None
 
+# =========================
+# SCRIPT DOWNLOADER
+# =========================
+def download_scripts():
+    resp = requests.get(SCRIPTS_REPO_API)
+    resp.raise_for_status()
 
+    scripts = []
+    for item in resp.json():
+        if item["name"].endswith(".py"):
+            scripts.append(item["name"])
+            path = os.path.join(LOCAL_SCRIPTS_DIR, item["name"])
+            if not os.path.exists(path):
+                content = requests.get(item["download_url"]).text
+                with open(path, "w") as f:
+                    f.write(content)
 
-# ============================================================
-# HELPERS
-# ============================================================
+    return sorted(scripts)
 
-def list_loras_from_repo(repo_id: str):
-    try:
-        files = list_repo_files(repo_id)
-        return [f for f in files if f.endswith(".safetensors")]
-    except Exception as e:
-        log(f"❌ Failed to list LoRAs: {e}")
-        return []
 
-# ============================================================
-# IMAGE GENERATION (SAFE LORA LOGIC)
-# ============================================================
-@spaces.GPU()
-def generate_image(prompt, height, width, steps, seed, guidance_scale):
-    LOGS = []
-    print(prompt)
+SCRIPT_NAMES = download_scripts()
 
-    if pipe is None:
-        return None, [], "❌ Pipeline not initialized"
 
-    generator = torch.Generator().manual_seed(int(seed))
-    placeholder = Image.new("RGB", (width, height), (255, 255, 255))
-    previews = []
+# =========================
+# PIPELINE BUILDER
+# =========================
+def build_pipeline(selected_scripts):
+    global pipe
 
-    # ---- Always start clean ----
-    try:
-        pipe.unload_lora_weights()
-    except Exception:
-        pass
-
-    # ---- Load LoRA for this run only ----
-    if CURRENT_LORA_FILE:
-        try:
-            pipe.load_lora_weights(
-                CURRENT_LORA_REPO,
-                weight_name=CURRENT_LORA_FILE
-            )
-            LOGS.append(f"🧩 LoRA loaded: {CURRENT_LORA_FILE}")
-        except Exception as e:
-            LOGS.append(f"❌ LoRA load failed: {e}")
-
-    # ---- Preview steps (lightweight) ----
-    try:
-        num_previews = min(5, steps)
-        for i in range(num_previews):
-            out = pipe(
-                prompt=prompt,
-                height=height // 4,
-                width=width // 4,
-                num_inference_steps=i + 1,
-                guidance_scale=guidance_scale,
-                generator=generator,
-            )
-            img = out.images[0].resize((width, height))
-            previews.append(img)
-            yield None, previews, "\n".join(LOGS)
-    except Exception as e:
-        LOGS.append(f"⚠️ Preview failed: {e}")
+    log("Building Z-Image Turbo pipeline...")
+    log(f"Selected scripts: {selected_scripts}")
 
-    # ---- Final image ----
-    try:
-        out = pipe(
-            prompt=prompt,
-            height=height,
-            width=width,
-            num_inference_steps=steps,
-            guidance_scale=guidance_scale,
-            generator=generator,
-        )
-        final_img = out.images[0]
-        previews.append(final_img)
-        LOGS.append("✅ Image generated")
+    pipe = DiffusionPipeline.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=False,
+        attn_implementation="kernels-community/vllm-flash-attn3",
+    )
 
-        yield final_img, previews, "\n".join(LOGS)
+    pipe.to("cuda")
+    log("Pipeline loaded and moved to CUDA")
 
-    except Exception as e:
-        LOGS.append(f"❌ Generation failed: {e}")
-        yield placeholder, previews, "\n".join(LOGS)
-
-    finally:
-        # ---- CRITICAL: unload after run ----
-        try:
-            pipe.unload_lora_weights()
-            LOGS.append("🧹 LoRA unloaded")
-        except Exception:
-            pass
-
-# ============================================================
-# GRADIO UI
-# ============================================================
-css = """
-.gradio-container {
-    max-width: 100% !important;
-    padding: 16px 32px !important;
-}
-
-.section {
-    margin-bottom: 12px;
-}
-
-.generate-btn {
-    background: linear-gradient(90deg, #4b6cb7, #182848) !important;
-    color: white !important;
-    font-weight: 600;
-    height: 46px;
-    border-radius: 10px;
-}
-
-.secondary-btn {
-    height: 42px;
-    border-radius: 10px;
-}
-
-textarea, input {
-    border-radius: 10px !important;
-}
-"""
+    return "Pipeline ready ✅"
 
-with gr.Blocks(
-    title="Z-Image-Turbo (Runtime LoRA)",
-    css=css,
-) as demo:
 
-    gr.Markdown(
-        """
-        # 🎨 Z-Image-Turbo LORA
-        **Runtime LoRA · Safe Mode · Full-Width UI**
-        """
+# =========================
+# IMAGE GENERATION
+# =========================
+@spaces.GPU
+def generate_image(
+    prompt,
+    height,
+    width,
+    num_inference_steps,
+    seed,
+    randomize_seed,
+    num_images,
+):
+    global pipe
+
+    log_buffer.truncate(0)
+    log_buffer.seek(0)
+
+    if pipe is None:
+        raise RuntimeError("Pipeline not built yet")
+
+    log("=== NEW GENERATION REQUEST ===")
+    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),
+        width=int(width),
+        num_inference_steps=int(num_inference_steps),
+        guidance_scale=0.0,
+        generator=generator,
+        max_sequence_length=1024,
+        num_images_per_prompt=num_images,
+        output_type="pil",
     )
 
-    # ======================================================
-    # MAIN LAYOUT
-    # ======================================================
-    with gr.Row():
+    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}")
 
-        # ================= LEFT PANEL =================
-        with gr.Column(scale=5):
-
-            # -------- Prompt --------
-            prompt = gr.Textbox(
-                label="Prompt",
-                value="boat in ocean",
-                lines=4,
-                placeholder="Describe the image you want to generate…",
-            )
-
-            # -------- LoRA Controls (NEXT TO PROMPT) --------
-            gr.Markdown("### 🧩 LoRA Controls")
-
-            lora_repo = gr.Textbox(
-                label="LoRA Repository",
-                value=DEFAULT_LORA_REPO,
-                lines=2,
-                placeholder="username/repo (e.g. rahul7star/ZImageLora)",
-            )
-
-            lora_dropdown = gr.Dropdown(
-                label="LoRA File",
-                choices=[],
-                interactive=True,
-            )
-
-            with gr.Row():
-                refresh_btn = gr.Button("🔄 Refresh LoRA List", elem_classes="secondary-btn")
-                clear_lora_btn = gr.Button("❌ Clear LoRA", elem_classes="secondary-btn")
-
-            # -------- Generation Controls --------
-            gr.Markdown("### ⚙️ Generation Settings")
-
-            with gr.Row():
-                width = gr.Slider(256, 2048, value=1024, step=8, label="Width")
-                height = gr.Slider(256, 2048, value=1024, step=8, label="Height")
-
-            with gr.Row():
-                steps = gr.Slider(1, 50, value=20, step=1, label="Steps")
-                guidance = gr.Slider(0, 10, value=0.0, step=0.5, label="Guidance")
-                seed = gr.Number(value=42, label="Seed", precision=0)
-
-            run_btn = gr.Button("🚀 Generate Image", elem_classes="generate-btn")
-
-            logs_box = gr.Textbox(
-                label="Logs",
-                lines=10,
-                interactive=False,
-            )
-
-        # ================= RIGHT PANEL =================
-        with gr.Column(scale=7):
-
-            final_image = gr.Image(
-                label="Final Image",
-                height=520,
-            )
-
-            gallery = gr.Gallery(
-                label="Generation Steps",
-                columns=4,
-                height=260,
-            )
-
-    # ======================================================
-    # CALLBACKS
-    # ======================================================
-    def refresh_loras(repo):
-        files = list_loras_from_repo(repo)
-        return gr.update(
-            choices=files,
-            value=files[0] if files else None,
+    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.Row():
+        script_selector = gr.CheckboxGroup(
+            choices=SCRIPT_NAMES,
+            label="Select Z-Image Scripts",
         )
 
-    refresh_btn.click(
-        refresh_loras,
-        inputs=[lora_repo],
-        outputs=[lora_dropdown],
+    build_btn = gr.Button("Build Pipeline")
+    build_status = gr.Textbox(label="Status", interactive=False)
+
+    build_btn.click(
+        fn=build_pipeline,
+        inputs=[script_selector],
+        outputs=[build_status],
     )
 
-    def select_lora(lora_file, repo):
-        global CURRENT_LORA_FILE, CURRENT_LORA_REPO
-        CURRENT_LORA_FILE = lora_file
-        CURRENT_LORA_REPO = repo
-        return f"🧩 Selected LoRA: {lora_file}"
+    gr.Markdown("---")
 
-    lora_dropdown.change(
-        select_lora,
-        inputs=[lora_dropdown, lora_repo],
-        outputs=[logs_box],
-    )
+    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")
 
-    def clear_lora():
-        global CURRENT_LORA_FILE, CURRENT_LORA_REPO
-        CURRENT_LORA_FILE = None
-        CURRENT_LORA_REPO = None
-        try:
-            pipe.unload_lora_weights()
-        except Exception:
-            pass
-        return (
-            gr.update(value=None),
-            "🧹 LoRA cleared — base model will be used."
-        )
+    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")
 
-    clear_lora_btn.click(
-        clear_lora,
-        outputs=[lora_dropdown, logs_box],
-    )
+    with gr.Row():
+        seed = gr.Number(value=0, label="Seed")
+        random_seed = gr.Checkbox(label="Randomize Seed", value=True)
 
-    run_btn.click(
-        generate_image,
-        inputs=[prompt, height, width, steps, seed, guidance],
-        outputs=[final_image, gallery, logs_box],
-    )
+    run_btn = gr.Button("Generate")
 
+    gallery = gr.Gallery(label="Results", columns=3)
+    seed_out = gr.Number(label="Used Seed")
+    logs = gr.Textbox(label="Logs", lines=12)
 
+    run_btn.click(
+        fn=generate_image,
+        inputs=[
+            prompt,
+            height,
+            width,
+            steps,
+            seed,
+            random_seed,
+            images,
+        ],
+        outputs=[gallery, seed_out, logs],
+    )
 
+demo.queue()
 demo.launch()