Full test of ui arch compatability

2026-08-28 06:56:05 -06:00
parent 520d96aac3
commit c3bc8b0b4e
4 changed files with 105 additions and 4 deletions
--- 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) ----
--- 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 = (
--- 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
--- 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": [