Revert "Fix OOM by deferring model loading and AOT compilation to runtime"

This reverts commit 49afc981b928f696a30599430f66fe00bca60375.

app.py

@@ -24,11 +24,6 @@ from PIL import Image
 import os
 import gradio as gr
 
-# --- Lazy Loading State ---
-# Pipeline is loaded lazily on first inference to avoid memory issues during build
-_pipeline_instance = None
-_pipeline_initialized = False
-
 def turn_into_video(input_images, output_images, prompt, progress=gr.Progress(track_tqdm=True)):
     if not input_images or not output_images:
         raise gr.Error("Please generate an output image first.")
@@ -93,60 +88,31 @@ def use_history_as_input(evt: gr.SelectData):
         return gr.update(value=[evt.value])
     return gr.update()
 
-# --- Model Loading (Lazy) ---
+# --- Model Loading ---
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-def get_pipeline():
-    """
-    Lazily initialize the pipeline on first use.
-    This avoids loading ~40GB of models during build, which causes OOM.
-    The pipeline is loaded and compiled on the first inference request
-    when GPU resources are available via ZeroGPU.
-    """
-    global _pipeline_instance, _pipeline_initialized
-
-    if _pipeline_initialized:
-        return _pipeline_instance
-
-    print("=== Initializing pipeline (first request) ===")
-
-    # Load the pipeline with transformer
-    pipe = QwenImageEditPlusPipeline.from_pretrained(
-        "Qwen/Qwen-Image-Edit-2509",
-        transformer=QwenImageTransformer2DModel.from_pretrained(
-            "linoyts/Qwen-Image-Edit-Rapid-AIO",
-            subfolder='transformer',
-            torch_dtype=dtype,
-            device_map='cuda'
-        ),
-        torch_dtype=dtype
-    ).to(device)
-
-    # Load and fuse LoRA weights
-    pipe.load_lora_weights(
-        "lovis93/next-scene-qwen-image-lora-2509",
-        weight_name="next-scene_lora-v2-3000.safetensors",
-        adapter_name="next-scene"
-    )
-    pipe.set_adapters(["next-scene"], adapter_weights=[1.])
-    pipe.fuse_lora(adapter_names=["next-scene"], lora_scale=1.)
-    pipe.unload_lora_weights()
+pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2509", 
+                                                transformer= QwenImageTransformer2DModel.from_pretrained("linoyts/Qwen-Image-Edit-Rapid-AIO", 
+                                                                                                         subfolder='transformer',
+                                                                                                         torch_dtype=dtype,
+                                                                                                         device_map='cuda'),torch_dtype=dtype).to(device)
 
-    # Apply optimizations
-    pipe.transformer.__class__ = QwenImageTransformer2DModel
-    pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
-
-    # AOT compilation (already in GPU context from infer())
-    print("=== Starting AOT compilation ===")
-    optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt", already_in_gpu_context=True)
-    print("=== AOT compilation complete ===")
+pipe.load_lora_weights(
+        "lovis93/next-scene-qwen-image-lora-2509", 
+        weight_name="next-scene_lora-v2-3000.safetensors", adapter_name="next-scene"
+    )
+pipe.set_adapters(["next-scene"], adapter_weights=[1.])
+pipe.fuse_lora(adapter_names=["next-scene"], lora_scale=1.)
+pipe.unload_lora_weights()
 
-    _pipeline_instance = pipe
-    _pipeline_initialized = True
 
-    return pipe
+# Apply the same optimizations from the first version
+pipe.transformer.__class__ = QwenImageTransformer2DModel
+pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
 
+# --- Ahead-of-time compilation ---
+optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")
 
 # --- UI Constants and Helpers ---
 MAX_SEED = np.iinfo(np.int32).max
@@ -158,7 +124,7 @@ def use_output_as_input(output_images):
     return output_images
 
 # --- Main Inference Function (with hardcoded negative prompt) ---
-@spaces.GPU(duration=600)  # Increased duration for first-request initialization
+@spaces.GPU(duration=300)
 def infer(
     images,
     prompt,
@@ -174,9 +140,6 @@ def infer(
     """
     Generates an image using the local Qwen-Image diffusers pipeline.
     """
-    # Get or initialize the pipeline (lazy loading)
-    pipe = get_pipeline()
-
     # Hardcode the negative prompt as requested
     negative_prompt = " "
 

optimization.py

@@ -45,17 +45,11 @@ INDUCTOR_CONFIGS = {
 }
 
 
-def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, already_in_gpu_context: bool = False, **kwargs: P.kwargs):
-    """
-    Optimize the pipeline transformer with AOT compilation.
+def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
 
-    Args:
-        pipeline: The diffusion pipeline to optimize
-        already_in_gpu_context: If True, skip the @spaces.GPU decorator since we're
-            already running in a GPU context (e.g., called from within infer())
-    """
+    @spaces.GPU(duration=1500)
+    def compile_transformer():
 
-    def _compile_transformer_impl():
         with spaces.aoti_capture(pipeline.transformer) as call:
             pipeline(*args, **kwargs)
 
@@ -63,7 +57,7 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, already_in_gpu
         dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
 
         # quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
-
+        
         exported = torch.export.export(
             mod=pipeline.transformer,
             args=call.args,
@@ -73,14 +67,4 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, already_in_gpu
 
         return spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
 
-    if already_in_gpu_context:
-        # We're already in a GPU context, run directly
-        compiled = _compile_transformer_impl()
-    else:
-        # Need to allocate GPU for compilation
-        @spaces.GPU(duration=1500)
-        def compile_transformer():
-            return _compile_transformer_impl()
-        compiled = compile_transformer()
-
-    spaces.aoti_apply(compiled, pipeline.transformer)
+    spaces.aoti_apply(compile_transformer(), pipeline.transformer)