Update qwenimage/pipeline_qwenimage_edit_plus.py

qwenimage/pipeline_qwenimage_edit_plus.py

@@ -166,13 +166,46 @@ def calculate_dimensions(target_area: int, ratio: float, multiple: int = 32):
 # Optional: decoder VAE (Wan2x)
 _ALT_VAE_WAN2X = None
 
+# Track desired tiling state for the optional decoder VAE, so it stays consistent across lazy loads.
+_ALT_VAE_WAN2X_TILING_ENABLED = False
+
+
+def _set_vae_tiling(model: Any, enabled: bool) -> bool:
+    """
+    Best-effort tiling toggle for a VAE-like module.
+    Returns True if a tiling method existed and was called, False otherwise.
+    """
+    if model is None:
+        return False
+    try:
+        if enabled:
+            if hasattr(model, "enable_tiling"):
+                model.enable_tiling()
+                return True
+            if hasattr(model, "enable_vae_tiling"):
+                model.enable_vae_tiling()
+                return True
+        else:
+            if hasattr(model, "disable_tiling"):
+                model.disable_tiling()
+                return True
+            if hasattr(model, "disable_vae_tiling"):
+                model.disable_vae_tiling()
+                return True
+    except Exception as e:
+        # Don't hard-fail inference if tiling toggle fails for an alt decoder.
+        logger.warning(f"VAE tiling toggle failed on {type(model)}: {e}")
+        return False
+    return False
+
+
 
 def _get_wan2x_vae(device: torch.device, dtype: torch.dtype):
     """
     Decoder-only finetune that outputs 2x resolution via pixel-shuffle.
     Lazy-loaded so it doesn't impact startup unless used.
     """
-    global _ALT_VAE_WAN2X
+    global _ALT_VAE_WAN2X, _ALT_VAE_WAN2X_TILING_ENABLED
     if _ALT_VAE_WAN2X is None:
         from diffusers import AutoencoderKLWan
 
@@ -182,7 +215,15 @@ def _get_wan2x_vae(device: torch.device, dtype: torch.dtype):
             torch_dtype=dtype,
         )
         _ALT_VAE_WAN2X.eval()
+
+        # Apply last requested tiling immediately on first load (if supported).
+        _set_vae_tiling(_ALT_VAE_WAN2X, _ALT_VAE_WAN2X_TILING_ENABLED)
+
     _ALT_VAE_WAN2X = _ALT_VAE_WAN2X.to(device=device, dtype=dtype)
+
+    # Re-apply after moving to device, just in case.
+    _set_vae_tiling(_ALT_VAE_WAN2X, _ALT_VAE_WAN2X_TILING_ENABLED)
+
     return _ALT_VAE_WAN2X
 
 
@@ -220,6 +261,9 @@ class QwenImageEditPlusPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
         self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
         self.tokenizer_max_length = 1024
 
+
+        # Track tiling state (applies to both primary VAE and optional decoder VAE)
+        self._vae_tiling_enabled = False
         self.prompt_template_encode = (
             "<|im_start|>system\n"
             "Describe the key features of the input image (color, shape, size, texture, objects, background), "
@@ -230,6 +274,32 @@ class QwenImageEditPlusPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
         self.prompt_template_encode_start_idx = 64
         self.default_sample_size = 128
 
+
+        # ------------------------------------------------------------
+        # VAE tiling control (applies to both primary VAE and optional decoder VAE)
+        # ------------------------------------------------------------
+        # Expose a stable API so app.py can call pipe.enable_vae_tiling()/disable_vae_tiling()
+        # regardless of which decoder VAE is selected at runtime.
+        
+    def set_vae_tiling(self, enabled: bool) -> None:
+        global _ALT_VAE_WAN2X_TILING_ENABLED, _ALT_VAE_WAN2X
+
+        enabled = bool(enabled)
+        self._vae_tiling_enabled = enabled
+
+        # 1) Primary VAE (Qwen)
+        _set_vae_tiling(getattr(self, "vae", None), enabled)
+
+        # 2) Optional decoder VAE (Wan2x): store desired global state; apply now if already loaded.
+        _ALT_VAE_WAN2X_TILING_ENABLED = enabled
+        if _ALT_VAE_WAN2X is not None:
+            _set_vae_tiling(_ALT_VAE_WAN2X, enabled)
+
+    def enable_vae_tiling(self) -> None:
+        self.set_vae_tiling(True)
+
+    def disable_vae_tiling(self) -> None:
+        self.set_vae_tiling(False)
     # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._extract_masked_hidden
     def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
         bool_mask = mask.bool()