diffusion: support Qwen-Image-Layered (#15817)

Co-authored-by: yhyang201 <yhyang201@gmail.com>

python/sglang/multimodal_gen/configs/models/vaes/qwenimage.py

@@ -20,7 +20,7 @@ class QwenImageVAEArchConfig(VAEArchConfig):
     dropout: float = 0.0
 
     is_residual: bool = False
-    in_channels: int = 3
+    input_channels: int = 3
     out_channels: int = 3
     patch_size: int | None = None
     scale_factor_temporal: int = 4

python/sglang/multimodal_gen/configs/pipeline_configs/base.py

@@ -299,6 +299,9 @@ class PipelineConfig:
 
         return shape
 
+    def allow_set_num_frames(self):
+        return False
+
     def get_decode_scale_and_shift(self, device, dtype, vae):
         vae_arch_config = self.vae_config.arch_config
         scaling_factor = getattr(vae_arch_config, "scaling_factor", None)

python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py

@@ -12,6 +12,7 @@ from sglang.multimodal_gen.configs.models.vaes.qwenimage import QwenImageVAEConf
 from sglang.multimodal_gen.configs.pipeline_configs.base import (
     ImagePipelineConfig,
     ModelTaskType,
+    maybe_unpad_latents,
     shard_rotary_emb_for_sp,
 )
 from sglang.multimodal_gen.runtime.models.vision_utils import resize
@@ -466,3 +467,90 @@ class QwenImageEditPlusPipelineConfig(QwenImageEditPipelineConfig):
             "freqs_cis": ((img_cos, img_sin), (txt_cos, txt_sin)),
             "img_shapes": img_shapes,
         }
+
+
+@dataclass
+class QwenImageLayeredPipelineConfig(QwenImageEditPipelineConfig):
+    resolution: int = 640  # TODO: allow user to set resolution
+    vae_precision: str = "bf16"
+
+    def _prepare_edit_cond_kwargs(
+        self, batch, prompt_embeds, rotary_emb, device, dtype
+    ):
+        batch_size = batch.latents.shape[0]
+        assert batch_size == 1
+        height = batch.height
+        width = batch.width
+        image_size = batch.original_condition_image_size
+
+        vae_scale_factor = self.get_vae_scale_factor()
+
+        img_shapes = batch.img_shapes
+        txt_seq_lens = batch.txt_seq_lens
+
+        (img_cos, img_sin), (txt_cos, txt_sin) = (
+            QwenImageEditPlusPipelineConfig.get_freqs_cis(
+                img_shapes, txt_seq_lens, rotary_emb, device, dtype
+            )
+        )
+
+        # perform sp shard on noisy image tokens
+        noisy_img_seq_len = (
+            1 * (height // vae_scale_factor // 2) * (width // vae_scale_factor // 2)
+        )
+        noisy_img_cos = shard_rotary_emb_for_sp(img_cos[:noisy_img_seq_len, :])
+        noisy_img_sin = shard_rotary_emb_for_sp(img_sin[:noisy_img_seq_len, :])
+
+        # concat back the img_cos for input image (since it is not sp-shared later)
+        img_cos = torch.cat([noisy_img_cos, img_cos[noisy_img_seq_len:, :]], dim=0).to(
+            device=device
+        )
+        img_sin = torch.cat([noisy_img_sin, img_sin[noisy_img_seq_len:, :]], dim=0).to(
+            device=device
+        )
+
+        return {
+            "txt_seq_lens": txt_seq_lens,
+            "img_shapes": img_shapes,
+            "freqs_cis": ((img_cos, img_sin), (txt_cos, txt_sin)),
+            "additional_t_cond": torch.tensor([0], device=device, dtype=torch.long),
+        }
+
+    def _unpad_and_unpack_latents(self, latents, batch):
+        vae_scale_factor = self.vae_config.arch_config.vae_scale_factor
+        channels = self.dit_config.arch_config.in_channels
+        batch_size = latents.shape[0]
+        layers = batch.num_frames
+
+        height = 2 * (int(batch.height) // (vae_scale_factor * 2))
+        width = 2 * (int(batch.width) // (vae_scale_factor * 2))
+
+        latents = maybe_unpad_latents(latents, batch)
+        latents = latents.view(
+            batch_size, layers + 1, height // 2, width // 2, channels // 4, 2, 2
+        )
+        latents = latents.permute(0, 1, 4, 2, 5, 3, 6)
+
+        latents = latents.reshape(
+            batch_size, layers + 1, channels // (2 * 2), height, width
+        )
+        latents = latents.permute(0, 2, 1, 3, 4)  # (b, c, f, h, w)
+        return latents, batch_size, channels, height, width
+
+    def allow_set_num_frames(self):
+        return True
+
+    def post_denoising_loop(self, latents, batch):
+        # unpack latents for qwen-image
+        (
+            latents,
+            batch_size,
+            channels,
+            height,
+            width,
+        ) = self._unpad_and_unpack_latents(latents, batch)
+        b, c, f, h, w = latents.shape
+        latents = latents[:, :, 1:]  # remove the first frame as it is the origin input
+        latents = latents.permute(0, 2, 1, 3, 4).view(-1, c, 1, h, w)
+        # latents = latents.reshape(batch_size, channels // (2 * 2), 1, height, width)
+        return latents

python/sglang/multimodal_gen/configs/sample/qwenimage.py

@@ -21,3 +21,17 @@ class QwenImageEditPlusSamplingParams(QwenImageSamplingParams):
     guidance_scale: float = 4.0
     # true_cfg_scale: float = 4.0
     num_inference_steps: int = 40
+
+
+@dataclass
+class QwenImageLayeredSamplingParams(QwenImageSamplingParams):
+    # num_frames: int = 4
+    height: int = 640
+    width: int = 640
+    prompt: str = " "
+    negative_prompt: str = " "
+
+    guidance_scale: float = 4.0
+    num_inference_steps: int = 50
+    cfg_normalize: bool = True
+    use_en_prompt: bool = True

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -69,7 +69,7 @@ class DataType(Enum):
 
     def get_default_extension(self) -> str:
         if self == DataType.IMAGE:
-            return "jpg"
+            return "png"
         else:
             return "mp4"
 
@@ -279,8 +279,10 @@ class SamplingParams:
 
         if pipeline_config.task_type.is_image_gen():
             # settle num_frames
-            logger.debug(f"num_frames set to 1 for image generation model")
-            self.num_frames = 1
+            if not server_args.pipeline_config.allow_set_num_frames():
+                logger.debug(f"num_frames set to 1 for image generation model")
+                self.num_frames = 1
+
         elif self.adjust_frames:
             # NOTE: We must apply adjust_num_frames BEFORE the SP alignment logic below.
             # If we apply it after, adjust_num_frames might modify the frame count