Update app_lora1.py

app_lora1.py

@@ -162,29 +162,42 @@ def get_pipeline(script_name):
     log(f"🔧 Building pipeline from {script_name}")
     code = SCRIPT_CODE[script_name]
 
-    # Extract lines after from_pretrained for logging
+    # Extract all lines after from_pretrained
     pipe_lines = extract_pipe_lines(code)
+    pipe_code = "\n".join(pipe_lines)  # <-- combine into full block
 
-    # Safe namespace for exec
-    namespace = {"torch": torch, "ZImagePipeline": ZImagePipeline}
+    # Safe namespace
+    from diffusers.models.transformers.z_image_transformer_2d import ZImageTransformer2DModel
+    from diffusers.models.model_loading_utils import GGUFQuantizationConfig
 
-    pipe = None
-    for line in pipe_lines:
-        if line.strip():
-            log(f"• {line.strip()}")
-            exec(line, namespace)
-            pipe = namespace.get("pipe", pipe)
+    namespace = {
+        "torch": torch,
+        "ZImagePipeline": ZImagePipeline,
+        "ZImageTransformer2DModel": ZImageTransformer2DModel,
+        "GGUFQuantizationConfig": GGUFQuantizationConfig,
+    }
 
+    # Execute entire block at once
+    try:
+        exec(pipe_code, namespace)
+    except Exception as e:
+        log(f"❌ Pipeline build failed: {e}")
+        raise RuntimeError(f"Pipeline build failed for {script_name}") from e
+
+    pipe = namespace.get("pipe", None)
     if pipe is None:
         raise RuntimeError(f"{script_name} did not define `pipe`.")
 
-    
+    # Move to CUDA only if not device_map
+    if not hasattr(pipe, "hf_device_map"):
+        pipe = pipe.to("cuda")
 
     PIPELINES[script_name] = pipe
     log("✅ Pipeline ready")
     return pipe
 
 
+
 # =========================================================
 # IMAGE GENERATION
 # =========================================================