Refactor app.py to consolidate pipeline loading and adapter management into a single shared pipeline for improved efficiency and clarity

app.py

@@ -47,37 +47,21 @@ scheduler_config = {
 # Initialize scheduler with Lightning config
 scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
 
-# Load Stage 1 pipeline (Anime -> Base Body)
-pipe_stage1 = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2511",
+# Load single shared pipeline
+pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2511",
                                                  scheduler=scheduler,
                                                  torch_dtype=dtype).to(device)
+# Load all LoRAs but don't fuse yet
 # Load 4-step Lightning LoRA
-pipe_stage1.load_lora_weights(
+pipe.load_lora_weights(
         "lightx2v/Qwen-Image-Edit-2511-Lightning",
         weight_name="Qwen-Image-Edit-2511-Lightning-4steps-V1.0-bf16.safetensors",
         adapter_name="lightning"
 )
 # Load Stage 1 LoRA
-pipe_stage1.load_lora_weights(STAGE1_LORA_REPO, weight_name=STAGE1_LORA_WEIGHT, adapter_name="stage1")
-# Set both adapters with equal weights and fuse
-pipe_stage1.set_adapters(["lightning", "stage1"], adapter_weights=[1.0, 1.0])
-pipe_stage1.fuse_lora()
-
-# Load Stage 2 pipeline (Base Body -> Guide Body)
-pipe_stage2 = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2511",
-                                                 scheduler=scheduler,
-                                                 torch_dtype=dtype).to(device)
-# Load 4-step Lightning LoRA
-pipe_stage2.load_lora_weights(
-        "lightx2v/Qwen-Image-Edit-2511-Lightning",
-        weight_name="Qwen-Image-Edit-2511-Lightning-4steps-V1.0-bf16.safetensors",
-        adapter_name="lightning"
-)
+pipe.load_lora_weights(STAGE1_LORA_REPO, weight_name=STAGE1_LORA_WEIGHT, adapter_name="stage1")
 # Load Stage 2 LoRA
-pipe_stage2.load_lora_weights(STAGE2_LORA_REPO, weight_name=STAGE2_LORA_WEIGHT, adapter_name="stage2")
-# Set both adapters with equal weights and fuse
-pipe_stage2.set_adapters(["lightning", "stage2"], adapter_weights=[1.0, 1.0])
-pipe_stage2.fuse_lora()
+pipe.load_lora_weights(STAGE2_LORA_REPO, weight_name=STAGE2_LORA_WEIGHT, adapter_name="stage2")
 
 # # Apply the same optimizations from the first version
 # pipe.transformer.__class__ = QwenImageTransformer2DModel
@@ -149,7 +133,10 @@ def infer(
     print(f"Prompt: '{STAGE1_PROMPT}'")
     print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}, Size: {width}x{height}")
 
-    stage1_images = pipe_stage1(
+    # Set Stage 1 adapters
+    pipe.set_adapters(["lightning", "stage1"], adapter_weights=[1.0, 1.0])
+
+    stage1_images = pipe(
         image=pil_images if len(pil_images) > 0 else None,
         prompt=STAGE1_PROMPT,
         height=height,
@@ -165,10 +152,13 @@ def infer(
     print(f"[Stage 2] Converting to guide body...")
     print(f"Prompt: '{STAGE2_PROMPT}'")
 
+    # Set Stage 2 adapters
+    pipe.set_adapters(["lightning", "stage2"], adapter_weights=[1.0, 1.0])
+
     # Use same seed for stage 2
     generator = torch.Generator(device=device).manual_seed(seed)
 
-    stage2_images = pipe_stage2(
+    stage2_images = pipe(
         image=stage1_images,
         prompt=STAGE2_PROMPT,
         height=height,