[New model] Qwen-Image-2.1 support (rebase + review follow-ups) (#8099)

Signed-off-by: NumberWan <wantszkin2003@gmail.com>
Signed-off-by: SamitHuang <285365963@qq.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Tianyu Guo <guoty9@mail2.sysu.edu.cn>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: SamitHuang <285365963@qq.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from typing import Any
+
+import torch
+from diffusers.models.autoencoders.vae import DecoderOutput
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedOperator,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+from vllm_omni.diffusion.models.qwen_image_21.autoencoder_kl_qwenimage21 import (
+    AutoencoderKLQwenImage21,
+    _cudnn_deterministic,
+    _unpatchify,
+)
+
+logger = init_logger(__name__)
+
+
+class DistributedAutoencoderKLQwenImage21(AutoencoderKLQwenImage21, DistributedVaeMixin):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any):
+        model = super().from_pretrained(*args, **kwargs)
+        model.init_distributed()
+        return model
+
+    def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+        # mostly copy from AutoencoderKLQwenImage21.tiled_decode
+        _, _, num_frames, height, width = z.shape
+        sample_height = height * self.spatial_compression_ratio
+        sample_width = width * self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        tile_sample_stride_height = self.tile_sample_stride_height
+        tile_sample_stride_width = self.tile_sample_stride_width
+        if self.config.patch_size is not None:
+            sample_height = sample_height // self.config.patch_size
+            sample_width = sample_width // self.config.patch_size
+            tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size
+            tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size
+            blend_height = self.tile_sample_min_height // self.config.patch_size - tile_sample_stride_height
+            blend_width = self.tile_sample_min_width // self.config.patch_size - tile_sample_stride_width
+        else:
+            blend_height = self.tile_sample_min_height - tile_sample_stride_height
+            blend_width = self.tile_sample_min_width - tile_sample_stride_width
+
+        # Split z into overlapping tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        tiletask_list: list[TileTask] = []
+        for i in range(0, height, tile_latent_stride_height):
+            for j in range(0, width, tile_latent_stride_width):
+                time_list = []
+                for k in range(num_frames):
+                    tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
+                    time_list.append(tile)
+                tiletask_list.append(
+                    TileTask(
+                        len(tiletask_list),
+                        (i // tile_latent_stride_height, j // tile_latent_stride_width),
+                        time_list,
+                        workload=time_list[0].shape[3] * time_list[0].shape[4],
+                    )
+                )
+        tile_spec = {
+            "sample_height": sample_height,
+            "sample_width": sample_width,
+            "blend_height": blend_height,
+            "blend_width": blend_width,
+            "tile_sample_stride_height": tile_sample_stride_height,
+            "tile_sample_stride_width": tile_sample_stride_width,
+        }
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec=tile_spec,
+            output_dtype=self.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def tile_exec(self, task: TileTask) -> torch.Tensor:
+        """Decode a single latent tile into RGB space."""
+        self.clear_cache()
+        time = []
+        with _cudnn_deterministic():
+            for k in range(len(task.tensor)):
+                self._conv_idx = [0]
+                tile = self.post_quant_conv(task.tensor[k])
+                decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0))
+                time.append(decoded)
+        result = torch.cat(time, dim=2)
+        return result
+
+    def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
+        """Merge decoded tiles into a full image."""
+        grid_h, grid_w = grid_spec.grid_shape
+        self.clear_cache()
+
+        result_rows = []
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                tile = coord_tensor_map[(i, j)]
+                if i > 0:
+                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+                if j > 0:
+                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+                result_row.append(
+                    tile[
+                        :,
+                        :,
+                        :,
+                        : grid_spec.tile_spec["tile_sample_stride_height"],
+                        : grid_spec.tile_spec["tile_sample_stride_width"],
+                    ]
+                )
+            result_rows.append(torch.cat(result_row, dim=-1))
+        dec = torch.cat(result_rows, dim=3)[
+            :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
+        ]
+
+        if self.config.patch_size is not None:
+            dec = _unpatchify(dec, patch_size=self.config.patch_size)
+
+        dec = torch.clamp(dec, min=-1.0, max=1.0)
+        return dec
+
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
+        if not self.is_distributed_enabled():
+            return super().tiled_decode(z, return_dict=return_dict)
+
+        logger.debug("Decode running with distributed executor")
+        result = self.distributed_executor.execute(
+            z,
+            DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
+            broadcast_result=True,
+        )
+        if not return_dict:
+            return (result,)
+
+        return DecoderOutput(sample=result)

vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py

@@ -3,6 +3,8 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 
 import torch
@@ -34,9 +36,9 @@ class TileTask:
 
 @dataclass
 class DistributedOperator:
-    split: callable
-    exec: callable
-    merge: callable
+    split: Callable[[torch.Tensor], tuple[list[TileTask], GridSpec]]
+    exec: Callable[[TileTask], torch.Tensor]
+    merge: Callable[[dict[tuple[int, ...], torch.Tensor], GridSpec], torch.Tensor]
 
 
 class DistributedVaeExecutor:
@@ -46,7 +48,11 @@ class DistributedVaeExecutor:
 
     def __init__(self):
         # Use a dedicated process group spanning the complete worker WORLD.
-        self.group = get_world_group().device_group
+        world_group = get_world_group()
+        self.group = world_group.device_group
+        self._error_group = world_group.cpu_group
+        # Reporting an OOM must not require another device allocation.
+        self._error_status = torch.zeros((), dtype=torch.int32, device="cpu")
         self.world_size = dist.get_world_size(self.group)
         self.rank = dist.get_rank(self.group)
         self.parallel_size = 1
@@ -56,8 +62,34 @@ class DistributedVaeExecutor:
         self.parallel_size = parallel_size
         self.parallel_mode = mode
 
+    @contextmanager
+    def _sync_errors(self, stage: str):
+        """Agree on local-stage failures before entering the next collective."""
+        status = 0
+        message = ""
+        try:
+            yield
+        except torch.OutOfMemoryError as error:
+            status = 1
+            message = str(error)
+        except Exception as error:
+            # A non-retryable error takes precedence over OOM on another rank.
+            status = 2
+            message = f"{type(error).__name__}: {error}"
+
+        self._error_status.fill_(status)
+        dist.all_reduce(self._error_status, op=dist.ReduceOp.MAX, group=self._error_group)
+        global_status = self._error_status.item()
+        if global_status == 2:
+            detail = message if status == 2 else "another rank failed"
+            raise RuntimeError(f"Distributed VAE {stage} failed: {detail}")
+        if global_status == 1:
+            detail = message if status == 1 else "another rank ran out of memory"
+            raise torch.OutOfMemoryError(f"Distributed VAE {stage} failed: {detail}")
+
     def gather_tensors(self, tensor: torch.Tensor):
-        gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)]
+        with self._sync_errors("gather allocation"):
+            gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)]
         dist.all_gather(gather_list, tensor, group=self.group)
         return gather_list if self.rank == 0 else None
 
@@ -66,11 +98,12 @@ class DistributedVaeExecutor:
         return tensor
 
     def _compute_global_padding_shape(self, local_results, output_ndim: int, device):
-        local_tile_max_dims = [0] * output_ndim
-        for _, tile_tensor in local_results:
-            for dim_idx, dim_size in enumerate(tile_tensor.shape):
-                local_tile_max_dims[dim_idx] = max(local_tile_max_dims[dim_idx], dim_size)
-        local_shape_stat = torch.tensor([len(local_results), *local_tile_max_dims], device=device)
+        with self._sync_errors("shape preparation"):
+            local_tile_max_dims = [0] * output_ndim
+            for _, tile_tensor in local_results:
+                for dim_idx, dim_size in enumerate(tile_tensor.shape):
+                    local_tile_max_dims[dim_idx] = max(local_tile_max_dims[dim_idx], dim_size)
+            local_shape_stat = torch.tensor([len(local_results), *local_tile_max_dims], device=device)
         dist.all_reduce(
             local_shape_stat,
             op=dist.ReduceOp.MAX,
@@ -100,17 +133,16 @@ class DistributedVaeExecutor:
                 tid = int(meta_src[idx, 0])
                 if tid < 0:
                     continue
-                slices = [slice(None)] * tiles_src[idx].ndim
+                slices_list = [slice(None)] * tiles_src[idx].ndim
                 for i, dim in enumerate(grid_spec.split_dims):
-                    slices[dim] = slice(0, int(meta_src[idx, i + 1]))
-                slices = tuple(slices)
-                coord_tensor_map[tid_coord_map[tid]] = tiles_src[idx][slices]
+                    slices_list[dim] = slice(0, int(meta_src[idx, i + 1]))
+                coord_tensor_map[tid_coord_map[tid]] = tiles_src[idx][tuple(slices_list)]
 
         return coord_tensor_map
 
     def _balance_tasks(self, task_list, num_rank):
         workloads = [0] * num_rank
-        assigned = [[] for _ in range(num_rank)]
+        assigned: list[list[TileTask]] = [[] for _ in range(num_rank)]
 
         for task in sorted(task_list, key=lambda t: t.workload, reverse=True):
             r = workloads.index(min(workloads))
@@ -122,49 +154,51 @@ class DistributedVaeExecutor:
     def execute(self, z: torch.Tensor, operator: DistributedOperator, broadcast_result: bool = True):
         pp_size = min(self.parallel_size, self.world_size)
 
-        # 1. Split into tiles
-        tiletask_list, grid_spec = operator.split(z)
-        tid_coord_map = {task.tile_id: task.grid_coord for task in tiletask_list}
-
-        # 2. local decode
-        assigned = self._balance_tasks(tiletask_list, pp_size)
-        local_tasks = assigned[self.rank] if self.rank < pp_size else []
-        local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks]
+        with self._sync_errors("tile decode"):
+            tiletask_list, grid_spec = operator.split(z)
+            tid_coord_map = {task.tile_id: task.grid_coord for task in tiletask_list}
+            assigned = self._balance_tasks(tiletask_list, pp_size)
+            local_tasks = assigned[self.rank] if self.rank < pp_size else []
+            local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks]
 
         # 3. compute shape per rank
         global_padding_shape = self._compute_global_padding_shape(local_results, z.ndim, z.device)
 
         # 5. prepare tile tensors
         output_dtype = grid_spec.output_dtype if grid_spec.output_dtype is not None else z.dtype
-        local_tile_tensor, local_meta_tensor = self._pack_local_tiles(
-            local_results, global_padding_shape, grid_spec, z.device, output_dtype
-        )
+        with self._sync_errors("tile packing"):
+            local_tile_tensor, local_meta_tensor = self._pack_local_tiles(
+                local_results, global_padding_shape, grid_spec, z.device, output_dtype
+            )
 
         # 6. gather tiles & meta
         meta_gather = self.gather_tensors(local_meta_tensor)
         tile_gather = self.gather_tensors(local_tile_tensor)
 
-        if self.rank != 0:
-            result = torch.empty(0, device=z.device)  # Dummy return for non-zero ranks
-        else:
-            # 7. reconstruct full tensor (rank 0)
-            coord_tensor_map = self._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map)
-            result = operator.merge(coord_tensor_map, grid_spec)
+        with self._sync_errors("tile merge"):
+            if self.rank != 0:
+                result = torch.empty(0, device=z.device)  # Dummy return for non-zero ranks
+            else:
+                # 7. reconstruct full tensor (rank 0)
+                coord_tensor_map = self._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map)
+                result = operator.merge(coord_tensor_map, grid_spec)
 
         if broadcast_result:
             result = self._sync_final_result(result, z.ndim, z.device, output_dtype)
         return result
 
     def _sync_final_result(self, rank0_result, output_ndim, output_device, output_dtype):
-        shape_tensor = torch.empty((output_ndim,), device=output_device, dtype=torch.int64)
-        if self.rank == 0:
-            shape_tensor.copy_(torch.tensor(tuple(rank0_result.shape), device=output_device, dtype=torch.int64))
+        with self._sync_errors("output shape preparation"):
+            shape_tensor = torch.empty((output_ndim,), device=output_device, dtype=torch.int64)
+            if self.rank == 0:
+                shape_tensor.copy_(torch.tensor(tuple(rank0_result.shape), device=output_device, dtype=torch.int64))
         dist.broadcast(shape_tensor, src=0, group=self.group)
 
-        if self.rank != 0:
-            sync_result = torch.empty(tuple(shape_tensor.tolist()), device=output_device, dtype=output_dtype)
-        else:
-            sync_result = rank0_result
+        with self._sync_errors("output allocation"):
+            if self.rank != 0:
+                sync_result = torch.empty(tuple(shape_tensor.tolist()), device=output_device, dtype=output_dtype)
+            else:
+                sync_result = rank0_result
         dist.broadcast(sync_result, src=0, group=self.group)
         return sync_result