From 002be6725793b905c054f6e7d53c8d811ddb6f51 Mon Sep 17 00:00:00 2001 From: dahai80 <121743945@qq.com> Date: Sat, 8 Aug 2026 11:09:01 +0800 Subject: [PATCH] fix(wan2): staged VAE decode Stream(gpu, N) cross-thread error (#410 follow-up) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Phase-2 staged pipeline (load_text_encoder → encode_text → load_dit → denoise → load_vae → decode) raised "RuntimeError: There is no Stream(gpu, 4) in current thread" at the VAE decode mx.eval. Stages 1 (text encode) + 2 (DiT denoise) passed; stage 3 failed. The monolith generate() path was unaffected because it runs encode/denoise/decode in one continuous executor call. Root cause: MLX Metal streams are thread-local. The staged path round-trips the denoised latent through the event-loop main thread between the denoise executor call and the decode executor call. A *lazy* mx array (or a [None]/[0] projection built off-thread) records the producing call's auto-allocated internal Stream(gpu, N); evaluating it on a different thread (or after that stream's table entry is gone) raises the error. Controlled experiments confirmed: a latent built inside the decode executor decodes fine; one built on the main thread fails; a concrete mx.array(numpy) is portable. Fix (2 files, +75/-17): - wan2.py denoise(): build the 5D [None] batch projection and mx.eval it on the executor thread before returning (was a main-thread lazy projection). An eval'd array is portable across threads. - wan2.py decode()/decode_tiled(): slice latent[0] inside _decode on the executor thread, not on the main thread. Fixed t-axis index (2 for 5D, 1 for 4D). - wan2.py load_vae(): get_executor("io") → get_executor("video") for consistency with decode (io is a different thread pool). - wan2.py _clear_mlx_cache(): 3 unload methods route mx.synchronize()/ mx.clear_cache() through the video executor thread (not main thread). - stage.py run_denoise(): mx.eval(latents) before return. - stage.py decode_wan_vae(): mx.eval(latent) at entry. Verification: - Real-model e2e (Wan2.1-T2V-1.3B, fusion-comfyui tests/e2e_wan2_staged.py): full staged T2V PASSES, 57-65s, output (1,20,480,832,3), non-zero pixel fraction 0.99. Sequential offload confirmed (mem 10.8GB→5.4GB→290MB). - Unit (tests/unit/test_wan2_stage_api.py): 20/20 pass. - ruff clean. Co-Authored-By: Claude --- fusion_mlx/engines/video_backends/wan2.py | 74 +++++++++++++++++------ fusion_mlx/video/wan2/stage.py | 18 ++++++ 2 files changed, 75 insertions(+), 17 deletions(-) diff --git a/fusion_mlx/engines/video_backends/wan2.py b/fusion_mlx/engines/video_backends/wan2.py index c8c2e26..3b1692d 100644 --- a/fusion_mlx/engines/video_backends/wan2.py +++ b/fusion_mlx/engines/video_backends/wan2.py @@ -36,6 +36,20 @@ def _active_mem() -> int: return -1 +async def _clear_mlx_cache() -> None: + # Run mx.synchronize()/mx.clear_cache() on the video executor thread, NOT + # the event-loop main thread. MLX Metal streams are thread-local: a + # clear_cache() issued from the main thread invalidates the worker's + # thread-local stream table, so the next run_in_executor call (e.g. VAE + # decode after unload_dit) raises "There is no Stream(gpu, N) in current + # thread". Keeping the sync+clear on the same worker thread that owns the + # streams preserves the table across staged unload/load boundaries (#410). + loop = asyncio.get_running_loop() + await loop.run_in_executor( + get_executor("video"), lambda: (mx.synchronize(), mx.clear_cache()) + ) + + # Max T5 text-embedding cache entries (LRU eviction when exceeded). _T5_EMBED_CACHE_MAX = 16 # Timeout for T5 encoder preload during start() — large model may take minutes. @@ -376,8 +390,7 @@ async def unload_text_encoder(self) -> None: self._t5_encoder = None self._stage_flags["text_encoder"] = False gc.collect() - mx.synchronize() - mx.clear_cache() + await _clear_mlx_cache() logger.info("stage:text_encoder unload wan2") async def load_dit(self) -> None: @@ -465,7 +478,7 @@ async def denoise( on_step = self._stage_on_step def _denoise(): - return run_denoise( + lat_4d = run_denoise( config, models, pos_embed, @@ -480,11 +493,23 @@ def _denoise(): no_compile, on_step=on_step, ) + # 5D contract: add batch dim -> (1, z_dim, t_latent, h_lat, w_lat). + # Build the projection AND evaluate it on THIS executor thread so + # the returned array is concrete, not a lazy graph that references + # this call's auto-allocated Stream(gpu, N). The staged path later + # runs VAE decode in a *separate* executor call and round-trips this + # array through the event-loop main thread; MLX Metal streams are + # thread-local, so a lazy array (or a main-thread [None] projection + # of one) built on this call's streams raises + # "There is no Stream(gpu, N) in current thread" at the decode-side + # mx.eval. An mx.eval'd array is portable across threads. The + # monolith shares one executor call so this is a no-op for it. + lat_5d = lat_4d[None] + mx.eval(lat_5d) + return lat_5d loop = asyncio.get_running_loop() - result_4d = await loop.run_in_executor(get_executor("video"), _denoise) - # 5D contract: add batch dim -> (1, z_dim, t_latent, h_lat, w_lat). - result = result_4d[None] + result = await loop.run_in_executor(get_executor("video"), _denoise) logger.info( "stage:dit denoise wan2 steps=%d cfg=%.2f out_shape=%s", steps, @@ -497,8 +522,7 @@ async def unload_dit(self) -> None: self._stage_dit_models = None self._stage_flags["dit"] = False gc.collect() - mx.synchronize() - mx.clear_cache() + await _clear_mlx_cache() logger.info("stage:dit unload wan2") async def load_vae(self) -> None: @@ -514,8 +538,14 @@ def _load(): return load_vae_decoder(vae_path, config) loop = asyncio.get_running_loop() + # Load VAE on the *video* executor (not "io"): MLX Metal streams are + # thread-local, and decode() runs on get_executor("video"). Weights + # loaded on a different (io) thread bind to that thread's streams; the + # decode-side mx.eval then raises "There is no Stream(gpu, N) in + # current thread". Matches load_dit() + denoise() both on "video", and + # the monolith generate() which load+decode on one executor call. self._stage_vae = await asyncio.wait_for( - loop.run_in_executor(get_executor("io"), _load), + loop.run_in_executor(get_executor("video"), _load), timeout=_T5_PRELOAD_TIMEOUT, ) self._stage_flags["vae"] = True @@ -529,10 +559,16 @@ async def decode(self, latent: mx.array) -> mx.array: raise RuntimeError("vae is unloaded; call load_vae().") config = self._ensure_stage_config() # Accept 5D (1, c, t, h, w) or 4D (c, t, h, w); run_denoise returns 5D. - lat_4d = latent[0] if latent.ndim == 5 else latent + # NOTE: do NOT slice latent[0] here on the main thread — that builds a + # lazy projection referencing the source array's (possibly cross-thread) + # stream, which raises "There is no Stream(gpu, N) in current thread" + # at decode-side mx.eval. Pass the full latent and slice on the executor + # thread inside _decode (decode_wan_vae then mx.eval's it locally). vae = self._stage_vae + ndim = latent.ndim def _decode(): + lat_4d = latent[0] if ndim == 5 else latent return decode_wan_vae(lat_4d, config, vae, tiling_config=None) loop = asyncio.get_running_loop() @@ -550,15 +586,20 @@ async def decode_tiled(self, latent: mx.array, tile_size: int = 256) -> mx.array if self._stage_vae is None: raise RuntimeError("vae is unloaded; call load_vae().") config = self._ensure_stage_config() - lat_4d = latent[0] if latent.ndim == 5 else latent + # Slice on the executor thread (see decode() note): main-thread + # latent[0] builds a lazy cross-thread projection. + vae = self._stage_vae + ndim = latent.ndim # tile_size is in pixels (ComfyUI convention); auto derives spatial+temporal. - height = lat_4d.shape[-2] * config.vae_stride[1] - width = lat_4d.shape[-1] * config.vae_stride[2] - num_frames = lat_4d.shape[1] * config.vae_stride[0] - 1 + # t-axis index: 2 for 5D (1,c,t,h,w), 1 for 4D (c,t,h,w). + height = latent.shape[-2] * config.vae_stride[1] + width = latent.shape[-1] * config.vae_stride[2] + t_idx = 2 if ndim == 5 else 1 + num_frames = latent.shape[t_idx] * config.vae_stride[0] - 1 tiling_config = TilingConfig.auto(height, width, num_frames) - vae = self._stage_vae def _decode(): + lat_4d = latent[0] if ndim == 5 else latent return decode_wan_vae(lat_4d, config, vae, tiling_config=tiling_config) loop = asyncio.get_running_loop() @@ -575,8 +616,7 @@ async def unload_vae(self) -> None: self._stage_vae = None self._stage_flags["vae"] = False gc.collect() - mx.synchronize() - mx.clear_cache() + await _clear_mlx_cache() logger.info("stage:vae unload wan2") def set_progress_callback(self, cb): diff --git a/fusion_mlx/video/wan2/stage.py b/fusion_mlx/video/wan2/stage.py index 724cd05..44b9211 100644 --- a/fusion_mlx/video/wan2/stage.py +++ b/fusion_mlx/video/wan2/stage.py @@ -437,6 +437,15 @@ def run_denoise( del model, kv gc.collect() mx.clear_cache() + # Ensure the returned latents are fully materialized on this executor + # thread. The staged path runs VAE decode in a *separate* executor call and + # round-trips this array through the event-loop main thread; MLX Metal + # streams are thread-local, so a still-lazy array evaluated later on the + # decode thread raises "There is no Stream(gpu, N) in current thread". The + # caller (wan2.py denoise()) further evaluates the batch-dim projection on + # this same thread before returning. The monolith shares one executor call + # so this is a no-op for it (still correct). + mx.eval(latents) return latents @@ -444,6 +453,15 @@ def decode_wan_vae(latent, config, vae, tiling_config=None): # VAE decode extracted from generate_video() lines 1149-1254 (T2V branch, # no I2V mask_blend). latent is 4D (z_dim, t_latent, h_lat, w_lat). # Returns uint8 frames [T, H, W, 3]. + # Materialize the incoming latent on this executor thread. In the staged + # path it is produced by denoise() in a *separate* executor call and + # round-trips through the event-loop main thread; MLX Metal streams are + # thread-local, so a still-lazy array (or a slice built off-thread) raises + # "There is no Stream(gpu, N) in current thread" at the decode-side + # mx.eval. Evaluating here on the decode thread makes it concrete and + # portable. The caller (wan2.py decode()) already slices on this thread, + # so the latent is local; this eval is the materialization guarantee. + mx.eval(latent) is_wan22_vae = config.vae_z_dim == 48 if is_wan22_vae: from .vae22 import denormalize_latents