Support int8_convrot VAE (#15334)

comfy/ldm/minimax/vae.py

@@ -199,11 +199,11 @@ class RotaryEmbeddingND(nn.Module):
 
 class FeedForward(nn.Module):
     # Gated SiLU FFN.
-    def __init__(self, dim, mult=4, bias=True):
+    def __init__(self, dim, mult=4, bias=True, operations=ops):
         super().__init__()
         inner_dim = dim * mult
-        self.w1 = ops.Linear(dim, inner_dim * 2, bias=bias)
-        self.w2 = ops.Linear(inner_dim, dim, bias=bias)
+        self.w1 = operations.Linear(dim, inner_dim * 2, bias=bias)
+        self.w2 = operations.Linear(inner_dim, dim, bias=bias)
 
     def forward(self, x):
         gate, x = self.w1(x).chunk(2, dim=-1)
@@ -211,15 +211,15 @@ class FeedForward(nn.Module):
 
 
 class Attention(nn.Module):
-    def __init__(self, heads, dim_head, bias=True, eps=1e-5):
+    def __init__(self, heads, dim_head, bias=True, eps=1e-5, operations=ops):
         super().__init__()
         self.dim_head = dim_head
         self.heads = heads
         inner_dim = dim_head * heads
         self.norm_q = ops.RMSNorm(dim_head, eps=eps, elementwise_affine=False)
         self.norm_k = ops.RMSNorm(dim_head, eps=eps, elementwise_affine=False)
-        self.to_qkv = ops.Linear(inner_dim, inner_dim * 3, bias=bias)
-        self.to_out = ops.Linear(inner_dim, inner_dim, bias=bias)
+        self.to_qkv = operations.Linear(inner_dim, inner_dim * 3, bias=bias)
+        self.to_out = operations.Linear(inner_dim, inner_dim, bias=bias)
 
     def forward(self, x, rotary_pos_emb=None):
         batch_size, seq_len, _ = x.shape
@@ -242,14 +242,14 @@ class Attention(nn.Module):
 
 
 class TransformerBlock(nn.Module):
-    def __init__(self, heads, dim_head, bias=True, eps=1e-5):
+    def __init__(self, heads, dim_head, bias=True, eps=1e-5, operations=ops):
         super().__init__()
         dim = heads * dim_head
         self.norm1 = ops.RMSNorm(dim, elementwise_affine=True, eps=eps)
-        self.attn = Attention(heads=heads, dim_head=dim_head, bias=bias, eps=eps)
+        self.attn = Attention(heads=heads, dim_head=dim_head, bias=bias, eps=eps, operations=operations)
         self.scale1 = nn.Parameter(torch.empty(dim))
         self.norm2 = ops.RMSNorm(dim, elementwise_affine=True, eps=eps)
-        self.ff = FeedForward(dim=dim, bias=bias)
+        self.ff = FeedForward(dim=dim, bias=bias, operations=operations)
         self.scale2 = nn.Parameter(torch.empty(dim))
 
     def forward(self, x, rotary_pos_emb=None):
@@ -259,7 +259,7 @@ class TransformerBlock(nn.Module):
 
 class ViT3DDecoder(nn.Module):
     def __init__(self, patch_size=16, patch_size_t=4, in_channels=24, out_channels=3, num_layers=36, heads=32, dim_head=64, rope_theta=100.0,
-                 rope_dim_ratio=0.75, bias=True, eps=1e-5, num_register_tokens=4):
+                 rope_dim_ratio=0.75, bias=True, eps=1e-5, num_register_tokens=4, operations=ops):
         super().__init__()
         dim = heads * dim_head
         self.patch_size = patch_size
@@ -274,7 +274,7 @@ class ViT3DDecoder(nn.Module):
         self.register_buffer("mask_token", torch.empty(1, 1, dim))
 
         self.transformer_blocks = nn.ModuleList(
-            [TransformerBlock(heads=heads, dim_head=dim_head, bias=bias, eps=eps)
+            [TransformerBlock(heads=heads, dim_head=dim_head, bias=bias, eps=eps, operations=operations)
              for _ in range(num_layers)]
         )
 
@@ -337,6 +337,7 @@ class MiniMaxH3VideoVAE(nn.Module):
         tile_size=256,
         tile_overlap_min=64,
         tiling=True,
+        operations=ops,
     ):
         super().__init__()
         self.vae_ratio = int(math.prod(space_down))
@@ -372,6 +373,7 @@ class MiniMaxH3VideoVAE(nn.Module):
             patch_size_t=self.vae_ratio_t,
             in_channels=z_channels,
             out_channels=out_ch,
+            operations=operations,
         )
 
         self.register_buffer("latents_mean", torch.tensor(LATENTS_MEAN))

comfy/sd.py

@@ -940,7 +940,11 @@ class VAE:
                 if not comfy.memory_management.aimdo_enabled:
                     self.disable_offload = True
             elif "decoder.transformer_blocks.0.scale1" in sd and "encoder.down.5.block.0.conv1.weight" in sd:  # MiniMax H3 video VAE
-                self.first_stage_model = comfy.ldm.minimax.vae.MiniMaxH3VideoVAE()
+                minimax_ops = comfy.ops.disable_weight_init
+                minimax_quant = comfy.utils.detect_layer_quantization(sd, "")
+                if minimax_quant is not None:  # int8+convrot quantized decoder
+                    minimax_ops = comfy.ops.mixed_precision_ops(minimax_quant, dtype if dtype is not None else torch.float16)
+                self.first_stage_model = comfy.ldm.minimax.vae.MiniMaxH3VideoVAE(operations=minimax_ops)
                 self.latent_channels = 24
                 self.latent_dim = 3
                 # frames 17k+5 <-> latents 5k+2, 16x spatial