Create app_lora1.py

app_lora1.py

@@ -0,0 +1,201 @@
+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
+
+# ============================================================
+# CONFIG
+# ============================================================
+
+MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
+DEFAULT_LORA_REPO = "rahul7star/ZImageLora"
+
+DTYPE = torch.bfloat16
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ============================================================
+# GLOBAL STATE
+# ============================================================
+
+pipe = None
+CURRENT_LORA_REPO = None
+CURRENT_LORA_FILE = None
+
+# ============================================================
+# LOGGING
+# ============================================================
+
+def log(msg):
+    print(msg)
+    return msg
+
+# ============================================================
+# PIPELINE BUILD (ONCE)
+# ============================================================
+
+try:
+    pipe = ZImagePipeline.from_pretrained(
+        MODEL_ID,
+        torch_dtype=DTYPE,
+    )
+
+    if hasattr(pipe, "transformer") and hasattr(pipe.transformer, "set_attention_backend"):
+        pipe.transformer.set_attention_backend("_flash_3")
+
+    pipe.to(DEVICE)
+    log("✅ Pipeline built successfully")
+
+except Exception as e:
+    log("❌ Pipeline build failed")
+    log(traceback.format_exc())
+    pipe = None
+
+# ============================================================
+# HELPERS
+# ============================================================
+
+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)
+# ============================================================
+
+def generate_image(prompt, height, width, steps, seed, guidance_scale):
+    LOGS = []
+
+    if pipe is None:
+        return None, [], "❌ Pipeline not initialized"
+
+    generator = torch.Generator(DEVICE).manual_seed(int(seed))
+    placeholder = Image.new("RGB", (width, height), (255, 255, 255))
+    previews = []
+
+    # ---- 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}")
+
+    # ---- 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")
+
+        yield final_img, previews, "\n".join(LOGS)
+
+    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
+# ============================================================
+
+with gr.Blocks(title="Z-Image-Turbo (LoRA Safe)") as demo:
+    gr.Markdown("# 🎨 Z-Image-Turbo — Runtime LoRA (SAFE MODE)")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            prompt = gr.Textbox(label="Prompt", value="boat in ocean")
+            height = gr.Slider(256, 2048, value=1024, step=8, label="Height")
+            width = gr.Slider(256, 2048, value=1024, step=8, label="Width")
+            steps = gr.Slider(1, 50, value=20, step=1, label="Inference Steps")
+            guidance = gr.Slider(0, 10, value=0.0, step=0.5, label="Guidance Scale")
+            seed = gr.Number(value=42, label="Seed")
+            run_btn = gr.Button("🚀 Generate")
+
+        with gr.Column(scale=1):
+            final_image = gr.Image(label="Final Image")
+            gallery = gr.Gallery(label="Steps", columns=4, height=256)
+
+    gr.Markdown("## 🧩 LoRA Controls")
+
+    with gr.Row():
+        lora_repo = gr.Textbox(label="LoRA Repo", value=DEFAULT_LORA_REPO)
+        lora_dropdown = gr.Dropdown(label="LoRA File", choices=[])
+        refresh_btn = gr.Button("🔄 Refresh")
+
+    logs_box = gr.Textbox(label="Logs", lines=18)
+
+    # ---- Callbacks ----
+
+    def refresh_loras(repo):
+        files = list_loras_from_repo(repo)
+        return gr.update(choices=files, value=files[0] if files else None)
+
+    refresh_btn.click(refresh_loras, inputs=[lora_repo], outputs=[lora_dropdown])
+
+    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}"
+
+    lora_dropdown.change(
+        select_lora,
+        inputs=[lora_dropdown, lora_repo],
+        outputs=[logs_box],
+    )
+
+    run_btn.click(
+        generate_image,
+        inputs=[prompt, height, width, steps, seed, guidance],
+        outputs=[final_image, gallery, logs_box],
+    )
+
+demo.launch()