Conversation
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughAdds documentation and runtime handling for a new Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
examples/diffusers/quantization/utils.py (1)
87-87: Tighten WAN block regex to avoid unintended matches.On Line 87, the pattern duplicates
blocks.39and uses unescaped.inblocks.0\./blocks.1\./blocks.2\., which can match unintended names. This may disable quantizers for extra modules.Proposed regex cleanup
- r".*(patch_embedding|condition_embedder|proj_out|blocks.0\.|blocks.1\.|blocks.2\.|blocks.39|blocks.38|blocks.39).*" + r".*(patch_embedding|condition_embedder|proj_out|blocks\.(0|1|2|38|39)\.).*"🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/diffusers/quantization/utils.py` at line 87, The regex pattern string r".*(patch_embedding|condition_embedder|proj_out|blocks.0\.|blocks.1\.|blocks.2\.|blocks.39|blocks.38|blocks.39).*" is too loose and duplicates blocks.39; update the pattern in utils.py to escape the unescaped dots in the "blocks.N" parts and remove the duplicate entry, and consolidate the block indices (e.g., use a grouped alternative like blocks\.(?:0|1|2|38|39)) so only the intended module names are matched and quantizers aren't accidentally disabled.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@examples/diffusers/quantization/quantize.py`:
- Around line 313-318: The code forwards a raw merged_base_safetensor_path to
export_hf_checkpoint which can fail later; before setting kwargs and calling
export_hf_checkpoint, validate that merged_path exists and is a file (use
Path(merged_path).exists() / is_file() or os.path.exists), and if invalid log a
clear error via self.logger.error and raise a ValueError (or return/exit) so
export_hf_checkpoint is not called with a bad path; update the block around
model_config.model_type == ModelType.LTX2 that assigns
kwargs["merged_base_safetensor_path"] accordingly.
---
Nitpick comments:
In `@examples/diffusers/quantization/utils.py`:
- Line 87: The regex pattern string
r".*(patch_embedding|condition_embedder|proj_out|blocks.0\.|blocks.1\.|blocks.2\.|blocks.39|blocks.38|blocks.39).*"
is too loose and duplicates blocks.39; update the pattern in utils.py to escape
the unescaped dots in the "blocks.N" parts and remove the duplicate entry, and
consolidate the block indices (e.g., use a grouped alternative like
blocks\.(?:0|1|2|38|39)) so only the intended module names are matched and
quantizers aren't accidentally disabled.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 4384aed0-4ef8-468a-88ac-b7d010bfa15b
📒 Files selected for processing (4)
examples/diffusers/README.mdexamples/diffusers/quantization/pipeline_manager.pyexamples/diffusers/quantization/quantize.pyexamples/diffusers/quantization/utils.py
| if model_config and model_config.model_type == ModelType.LTX2: | ||
| merged_path = model_config.extra_params.get("merged_base_safetensor_path") | ||
| if merged_path: | ||
| self.logger.info(f"Merging base safetensors from {merged_path} for LTX2 export") | ||
| kwargs["merged_base_safetensor_path"] = merged_path | ||
| export_hf_checkpoint(pipe, export_dir=self.config.hf_ckpt_dir, **kwargs) |
There was a problem hiding this comment.
Validate merged_base_safetensor_path before forwarding to export.
Current flow forwards raw user-provided value; a bad path will fail deeper in export. Adding a local existence check gives a clearer error and faster fail.
Proposed validation guard
if model_config and model_config.model_type == ModelType.LTX2:
merged_path = model_config.extra_params.get("merged_base_safetensor_path")
if merged_path:
- self.logger.info(f"Merging base safetensors from {merged_path} for LTX2 export")
- kwargs["merged_base_safetensor_path"] = merged_path
+ merged_path = Path(merged_path)
+ if not merged_path.is_file():
+ raise FileNotFoundError(
+ f"merged_base_safetensor_path does not exist or is not a file: {merged_path}"
+ )
+ self.logger.info(f"Merging base safetensors from {merged_path} for LTX2 export")
+ kwargs["merged_base_safetensor_path"] = str(merged_path)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@examples/diffusers/quantization/quantize.py` around lines 313 - 318, The code
forwards a raw merged_base_safetensor_path to export_hf_checkpoint which can
fail later; before setting kwargs and calling export_hf_checkpoint, validate
that merged_path exists and is a file (use Path(merged_path).exists() /
is_file() or os.path.exists), and if invalid log a clear error via
self.logger.error and raise a ValueError (or return/exit) so
export_hf_checkpoint is not called with a bad path; update the block around
model_config.model_type == ModelType.LTX2 that assigns
kwargs["merged_base_safetensor_path"] accordingly.
kevalmorabia97
requested review from
Edwardf0t1
and removed request for
kevalmorabia97
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #972 +/- ##
==========================================
- Coverage 72.13% 72.10% -0.03%
==========================================
Files 209 209
Lines 23631 23631
==========================================
- Hits 17046 17040 -6
- Misses 6585 6591 +6 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Pull request overview
Adds support for exporting an LTX-2 quantized checkpoint in a ComfyUI-compatible single-file safetensors format by plumbing a new merged_base_safetensor_path option through the diffusers quantization example workflow.
Changes:
- Extend HF checkpoint export in
quantize.pyto pass LTX2-specific export kwargs (merged_base_safetensor_path). - Ignore
merged_base_safetensor_pathwhen constructing the LTX2 pipeline so it doesn’t get forwarded as a pipeline kwarg. - Update WAN model quantization filter pattern and document the new export parameters in the diffusers README.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| examples/diffusers/quantization/utils.py | Updates WAN filter regex used to disable/select quantizers for specific module name patterns. |
| examples/diffusers/quantization/quantize.py | Adds LTX2-specific HF export behavior to enable ComfyUI-style merged safetensors export. |
| examples/diffusers/quantization/pipeline_manager.py | Drops the new export-only extra param from LTX2 pipeline construction kwargs. |
| examples/diffusers/README.md | Documents LTX-2 HF export directory and merged-base safetensors flag in the example command. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| --quantized-torch-ckpt-save-path ./ltx-2-transformer.pt | ||
| --quantized-torch-ckpt-save-path ./ltx-2-transformer.pt \ | ||
| --hf-ckpt-dir ./LTX2-NVFP4/ \ | ||
| --extra-param merged_base_safetensor_path=./ltx-2-19b-dev-fp8.safetensors |
There was a problem hiding this comment.
This section header says “FP4 (torch checkpoint export)”, but the example now also documents HuggingFace export (--hf-ckpt-dir) and ComfyUI-style merging (merged_base_safetensor_path). Consider updating the heading/nearby text to reflect that this example covers HF/ComfyUI checkpoint export as well, so readers aren’t confused about what gets produced.
Edwardf0t1
left a comment
Edwardf0t1
left a comment
There was a problem hiding this comment.
LGTM. Please check copilot's reviews as well.
Thank you. I’ve addressed it, and I found their suggestion genuinely useful. |
What does this PR do?
Type of change: Bug fix
Usage
python quantize.py \ --model ltx-2 --format fp4 --batch-size 1 --calib-size 32 --n-steps 40 \ --extra-param checkpoint_path=./ltx-2-19b-dev-fp8.safetensors \ --extra-param distilled_lora_path=./ltx-2-19b-distilled-lora-384.safetensors \ --extra-param spatial_upsampler_path=./ltx-2-spatial-upscaler-x2-1.0.safetensors \ --extra-param gemma_root=./gemma-3-12b-it-qat-q4_0-unquantized \ --extra-param fp8transformer=true \ --quantized-torch-ckpt-save-path ./ltx-2-transformer.pt \ --hf-ckpt-dir ./LTX2-NVFP4/ \ --extra-param merged_base_safetensor_path=./ltx-2-19b-dev-fp8.safetensorsTesting
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True, usingtorch.load(..., weights_only=True), avoidingpickle, etc.).Additional Information
Summary by CodeRabbit
Documentation
Improvements