diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py index 620254cf9..753050c41 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3.py +++ b/tests/e2e/accuracy/test_hunyuan_image3.py @@ -71,6 +71,7 @@ def _assert_cot_structural_markers(cot_text: str, label: str) -> None: os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA" +os.environ["VLLM_LOGGING_LEVEL"] = "DEBUG" pytestmark = [pytest.mark.local_model, pytest.mark.diffusion] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index 20ed54b4f..5ce3098ad 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -1065,9 +1065,10 @@ class ImageKVCacheManager: assert self.image_kv_cache_lens is not None if position_ids is not None: assert position_ids.shape == (bs, q_len) - assert torch.all(position_ids[:, 0] == self.image_kv_cache_lens), ( - "The first current position must immediately follow each sample's cached prompt KV." - ) + if logger.isEnabledFor(logging.DEBUG): + assert torch.all(position_ids[:, 0] == self.image_kv_cache_lens.to(position_ids.device)), ( + "The first current position must immediately follow each sample's cached prompt KV." + ) new_key = torch.cat([cached_key, key], dim=1) new_value = torch.cat([cached_value, value], dim=1) return new_key.contiguous(), new_value.contiguous() diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index 6635168b9..2aef06f20 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py @@ -987,7 +987,7 @@ class HunyuanImage3Pipeline( timestep_scatter_src = self.timestep_emb(t.reshape(-1)).reshape(batch_size, -1, n_embd) x.scatter_( dim=1, - index=timestep_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd), + index=timestep_scatter_index.to(x.device).unsqueeze(-1).repeat(1, 1, n_embd), src=timestep_scatter_src, ) @@ -1370,7 +1370,7 @@ class HunyuanImage3Pipeline( num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, image_mask=to_device(output.gen_image_mask, device), - gen_timestep_scatter_index=to_device(output.gen_timestep_scatter_index, device), + gen_timestep_scatter_index=output.gen_timestep_scatter_index, cond_vae_images=to_device(cond_vae_images, device), cond_timestep=to_device(cond_timestep, device), cond_vae_image_mask=to_device(output.cond_vae_image_mask, device),