Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for Zoe-style distillation in the Qwen image model, adding new configuration files, inference scripts, and scheduler logic for linear time shifting and Zoe-aligned latent generation. Feedback focuses on improving the robustness of latent preparation by avoiding fragile dimension squeezing and enhancing portability by replacing hardcoded absolute paths in scripts and configuration files with relative paths or environment variables.
| zoe_shape = (b, num_channels_latents * 4, t, height // 2, width // 2) | ||
| latents = randn_tensor(zoe_shape, generator=self.generator, device=AI_DEVICE, dtype=self.dtype) | ||
| # Convert to LightX2V sequence format: [B, (H//2)*(W//2), C*4] | ||
| latents = latents.squeeze(2) # [B, C*4, H//2, W//2] |
There was a problem hiding this comment.
Using squeeze(2) is fragile because it only removes the dimension if its size is exactly 1. If the temporal dimension t is greater than 1, this operation will do nothing, and the subsequent permute call will fail as it expects a 4D tensor but receives a 5D one. Since this is an image-specific latent preparation, it is safer to explicitly select the first frame to ensure the expected 4D shape [B, C*4, H//2, W//2].
| latents = latents.squeeze(2) # [B, C*4, H//2, W//2] | |
| latents = latents[:, :, 0, :, :] # [B, C*4, H//2, W//2] |
| "prompt_template_encode_start_idx": 34, | ||
| "attn_type": "sage_attn2", | ||
| "enable_cfg": false, | ||
| "dit_original_ckpt": "/data/nvme1/yongyang/ccc/models/distill_zoe_diff_qwen_image_data_680w_neo_prompt_res2k_3kiter_multi_large_char_200iter_step4.safetensors", |
There was a problem hiding this comment.
The path to the checkpoint is hardcoded to a specific absolute path (/data/nvme1/yongyang/...). This makes the configuration file non-portable and will cause failures in other environments. Consider using a relative path or a mechanism to resolve the model path dynamically.
| "attn_type": "sage_attn2", | ||
| "enable_cfg": false, | ||
| "dit_quantized": true, | ||
| "dit_quantized_ckpt": "/data/nvme1/yongyang/ccc/models/distill_zoe_diff_qwen_image_data_680w_neo_prompt_res2k_3kiter_multi_large_char_200iter_step4_fp8_mix.safetensors", |
There was a problem hiding this comment.
The path to the quantized checkpoint is hardcoded to a specific absolute path. This will cause the configuration to fail in any environment other than the one where it was created. Please use relative paths or environment variables.
| lightx2v_path=/data/nvme1/yongyang/ccc/LightX2V | ||
| model_path=/data/nvme1/models/Qwen/Qwen-Image-2512 |
There was a problem hiding this comment.
Hardcoded absolute paths for lightx2v_path and model_path limit the usability of this script across different systems. It is recommended to use relative paths or allow these paths to be passed as environment variables or command-line arguments.
| lightx2v_path=/data/nvme1/yongyang/ccc/LightX2V | ||
| model_path=/data/nvme1/models/Qwen/Qwen-Image-2512 |
There was a problem hiding this comment.
These absolute paths are specific to the author's environment. For better maintainability and sharing, consider using relative paths or environment variables.
2 participants
No description provided.