Fix device mismatch during AOT compilation warmup

Move text encoder offloading inside optimize_pipeline_() to happen
AFTER the warmup pass but BEFORE torch.export. This ensures:
1. Warmup run has all components on GPU (avoids CPU/CUDA mismatch)
2. torch.export has reduced memory (text encoder on CPU)
3. Text encoder returns to GPU after compilation

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

app.py

@@ -1,4 +1,3 @@
-import gc
 import gradio as gr
 import numpy as np
 import random
@@ -123,16 +122,9 @@ pipe.transformer.__class__ = QwenImageTransformer2DModel
 pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
 
 # --- Ahead-of-time compilation ---
-# Temporarily offload text encoder to CPU to free ~16GB during memory-intensive torch.export
-pipe.text_encoder.to('cpu')
-gc.collect()
-torch.cuda.empty_cache()
-
+# Note: optimize_pipeline_ handles text encoder offloading internally to save memory during torch.export
 optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")
 
-# Move text encoder back to GPU for inference
-pipe.text_encoder.to(device)
-
 # --- UI Constants and Helpers ---
 MAX_SEED = np.iinfo(np.int32).max
 

optimization.py

@@ -1,6 +1,8 @@
 """
+AOT compilation optimization for Qwen-Image-Edit pipeline.
 """
 
+import gc
 from typing import Any
 from typing import Callable
 from typing import ParamSpec
@@ -50,14 +52,22 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
     @spaces.GPU(duration=1500)
     def compile_transformer():
 
+        # Run warmup pass to capture transformer inputs (needs all components on GPU)
         with spaces.aoti_capture(pipeline.transformer) as call:
             pipeline(*args, **kwargs)
 
+        # Offload text encoder to CPU to free ~16GB during memory-intensive torch.export
+        # This happens AFTER warmup but BEFORE export to avoid device mismatch
+        text_encoder_device = next(pipeline.text_encoder.parameters()).device
+        pipeline.text_encoder.to('cpu')
+        gc.collect()
+        torch.cuda.empty_cache()
+
         dynamic_shapes = tree_map(lambda t: None, call.kwargs)
         dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
 
         # quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
-        
+
         exported = torch.export.export(
             mod=pipeline.transformer,
             args=call.args,
@@ -65,6 +75,11 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
             dynamic_shapes=dynamic_shapes,
         )
 
-        return spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
+        compiled = spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
+
+        # Move text encoder back to original device
+        pipeline.text_encoder.to(text_encoder_device)
+
+        return compiled
 
     spaces.aoti_apply(compile_transformer(), pipeline.transformer)