diff --git a/extensions_built_in/diffusion_models/ltx2/ltx2.py b/extensions_built_in/diffusion_models/ltx2/ltx2.py index ddd6f3f..94ca300 100644 --- a/extensions_built_in/diffusion_models/ltx2/ltx2.py +++ b/extensions_built_in/diffusion_models/ltx2/ltx2.py @@ -1374,8 +1374,12 @@ class LTX25Model(LTX2Model): transformer, transformer_sd, "transformer" ) del transformer_sd - if num_quantized_dit == 0: - transformer = transformer.to(dtype) + # cast to the compute dtype even when pre-quantized: the comfy file + # stores the scale_shift modulation tables in fp32, which promotes + # hidden states to fp32 and breaks the (bf16) diffusers attention + # linears. Quantized backends are immune (int8 qdata, uint8-viewed + # scales survive a dtype cast untouched). + transformer = transformer.to(dtype) flush() if self.model_config.quantize: @@ -1420,8 +1424,7 @@ class LTX25Model(LTX2Model): connectors, connectors_sd, "connectors" ) del connectors_sd, dit_sd - if num_quantized_connectors == 0: - connectors = connectors.to(dtype) + connectors = connectors.to(dtype) flush() # ---- text encoder (Gemma-4 12B, single comfy file) ---- diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index 9ad0882..77ae4ca 100644 --- a/testing/test_model_loading.py +++ b/testing/test_model_loading.py @@ -107,6 +107,77 @@ MODEL_TESTS = { "model": {"name_or_path": "black-forest-labs/FLUX.2-klein-base-4B", "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0}, }, + # ---- coverage for every UI-default arch (big downloads on first run) ---- + "wan22_14b": { + "model": {"name_or_path": "ai-toolkit/Wan2.2-T2V-A14B-Diffusers-bf16", "quantize": True, "quantize_te": True, "low_vram": True}, + "sample": {"width": 480, "height": 480, "num_inference_steps": 20, "guidance_scale": 3.5, "seed": 42, "num_frames": 17}, + }, + "wan22_14b_i2v": { + "model": {"name_or_path": "ai-toolkit/Wan2.2-I2V-A14B-Diffusers-bf16", "quantize": True, "quantize_te": True, "low_vram": True}, + "sample": {"width": 480, "height": 480, "num_inference_steps": 20, "guidance_scale": 3.5, "seed": 42, "num_frames": 17}, + "needs_control_image": True, + }, + "wan21_i2v": { + "model": {"name_or_path": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", "quantize": True, "quantize_te": True}, + "sample": {"width": 480, "height": 480, "num_inference_steps": 20, "guidance_scale": 5.0, "seed": 42, "num_frames": 17}, + "needs_control_image": True, + }, + "hidream": { + "model": {"name_or_path": "HiDream-ai/HiDream-I1-Full", "quantize": True, "quantize_te": True}, + "sample": {"width": 1024, "height": 1024, "num_inference_steps": 28, "guidance_scale": 5.0, "seed": 42}, + }, + "hidream_e1": { + # editing seq budget fits 768x768 (source+target concat) + "model": {"name_or_path": "HiDream-ai/HiDream-E1-1", "quantize": True, "quantize_te": True}, + "sample": {"width": 768, "height": 768, "num_inference_steps": 28, "guidance_scale": 5.0, "seed": 42}, + "needs_control_image": True, + }, + "nucleus_image": { + "model": {"name_or_path": "NucleusAI/Nucleus-Image", "quantize": True, "quantize_te": True}, + "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, + }, + "omnigen2": { + "model": {"name_or_path": "OmniGen2/OmniGen2", "quantize": True, "quantize_te": True}, + "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, + }, + "ltx2.5": { + "model": {"name_or_path": "Lightricks/LTX-2.5", "quantize": True, "quantize_te": True}, + "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 3.0, "seed": 42, "num_frames": 25}, + }, + "flux2": { + "model": {"name_or_path": "black-forest-labs/FLUX.2-dev", "quantize": True, "quantize_te": True}, + "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0}, + }, + "flux2_klein_9b": { + "model": {"name_or_path": "black-forest-labs/FLUX.2-klein-base-9B", "quantize": True, "quantize_te": True}, + "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0}, + }, + "prx_pixel": { + "model": {"name_or_path": "Photoroom/prxpixel-t2i", "quantize_te": True}, + "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0}, + }, + "zeta_chroma": { + "model": {"name_or_path": "lodestones/Zeta-Chroma/zeta-chroma-base-x0-pixel-dino-distance.safetensors", "extras_name_or_path": "Tongyi-MAI/Z-Image-Turbo", "quantize": True, "quantize_te": True}, + "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0}, + }, + "zimage_l2p": { + "model": {"name_or_path": "zhen-nan/L2P/model-1k-merge.safetensors", "extras_name_or_path": "Tongyi-MAI/Z-Image-Turbo", "quantize_te": True}, + "sample": {**IMG, "guidance_scale": 1.0}, + }, + "qwen_image_edit": { + "model": {"name_or_path": "Qwen/Qwen-Image-Edit", "quantize": True, "quantize_te": True}, + "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 4.0}, + "needs_control_image": True, + }, + "qwen_image_edit_plus": { + "model": {"name_or_path": "Qwen/Qwen-Image-Edit-2509", "quantize": True, "quantize_te": True}, + "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 4.0}, + "needs_control_image": True, + }, + "f-lite": { + "model": {"name_or_path": "Freepik/F-Lite", "quantize": True, "quantize_te": True}, + "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, + }, } SKIP_MARKERS = ( diff --git a/toolkit/models/v2/PLANNING.md b/toolkit/models/v2/PLANNING.md index 74096aa..3e943cc 100644 --- a/toolkit/models/v2/PLANNING.md +++ b/toolkit/models/v2/PLANNING.md @@ -408,6 +408,20 @@ Decisions: load → comfy save → identical key set + bit-exact quantized entries vs the published file → reload → bit-identical quantized forward. Extend per arch as saves flip. +- [x] Full UI-default coverage (2026-08-27): the registry now holds all 31 + UI-facing archs, and 30/31 load + generate through the new stack with + their real UI defaults (mageflow blocked upstream by its hub 404). + This round verified the previously-untested tail with real weights: + wan22_14b + i2v (comfy fp8 pairs via the legacy importer, candidate + keys added for the ai-toolkit bf16 default repos), wan21_i2v, hidream, + hidream_e1 (native 768² editing), nucleus, omnigen2, ltx2.5, flux2, + flux2_klein_9b, prx_pixel, zeta_chroma, zimage_l2p, both qwen edit + archs, f-lite. Fixes found by the run: ltx2.5's fp32 scale_shift + tables promoted hidden states into bf16 linears under the diffusers + class (ComfyUI casts per-op) — the DiT/connectors now cast to compute + dtype after quantized attach (ConvRot storage immune); harness configs + for e1 resolution and zeta/l2p extras_name_or_path corrected to match + the UI defaults. - [ ] Each newly migrated model adds its test in the same PR as its migration. ## TODO / look at later diff --git a/toolkit/models/v2/diffusion_models/wan.py b/toolkit/models/v2/diffusion_models/wan.py index dfa5958..67f62b6 100644 --- a/toolkit/models/v2/diffusion_models/wan.py +++ b/toolkit/models/v2/diffusion_models/wan.py @@ -30,6 +30,19 @@ class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin): ("Wan-AI/Wan2.2-I2V-A14B-Diffusers", "transformer_2"): [ "split_files/diffusion_models/wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors", ], + # the UI defaults point at the ai-toolkit bf16 repacks; same comfy files + ("ai-toolkit/Wan2.2-T2V-A14B-Diffusers-bf16", "transformer"): [ + "split_files/diffusion_models/wan2.2_t2v_high_noise_14B_fp8_scaled.safetensors", + ], + ("ai-toolkit/Wan2.2-T2V-A14B-Diffusers-bf16", "transformer_2"): [ + "split_files/diffusion_models/wan2.2_t2v_low_noise_14B_fp8_scaled.safetensors", + ], + ("ai-toolkit/Wan2.2-I2V-A14B-Diffusers-bf16", "transformer"): [ + "split_files/diffusion_models/wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors", + ], + ("ai-toolkit/Wan2.2-I2V-A14B-Diffusers-bf16", "transformer_2"): [ + "split_files/diffusion_models/wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors", + ], "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": { "repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged", "files": [