Add SeedVR2 7B support (#352)

Author: icelaglace

README.md

@@ -122,7 +122,7 @@ MFLUX supports the following model families. They have different strengths and w
 |[Z-Image](src/mflux/models/z_image/README.md) | Nov 2025 | 6B | Distilled & Base | Yes | Best all-rounder: fast, small, very good quality and realism. |
 |[FLUX.2](src/mflux/models/flux2/README.md) | Jan 2026 | 4B & 9B | Distilled & Base | Yes | Fastest + smallest with very good qaility and edit capabilities. |
 |[FIBO](src/mflux/models/fibo/README.md) | Oct 2025 | 8B | Base | No | Very good JSON-based prompt understanding and editability, medium speed |
-|[SeedVR2](src/mflux/models/seedvr2/README.md) | Jun 2025 | 3B | — | No | Best upscaling model. |
+|[SeedVR2](src/mflux/models/seedvr2/README.md) | Jun 2025 | 3B & 7B | — | No | Best upscaling model. |
 |[Qwen Image](src/mflux/models/qwen/README.md) | Aug 2025+ | 20B | Base | No | Large model (slower); strong prompt understanding and world knowledge. Has edit capabilities |
 |[Depth Pro](src/mflux/models/depth_pro/README.md) | Oct 2024 | — | — | No | Very fast and accurate depth estimation model from Apple. |
 |[FLUX.1](src/mflux/models/flux/README.md) | Aug 2024 | 12B | Distilled & Base | No (legacy) | Legacy option with decent quality. Has edit capabilities with 'Kontext' model and upscaling support via ControlNet |

src/mflux/models/common/config/model_config.py

@@ -142,6 +142,11 @@ def z_image() -> "ModelConfig":
     def seedvr2_3b() -> "ModelConfig":
         return AVAILABLE_MODELS["seedvr2-3b"]
 
+    @staticmethod
+    @lru_cache
+    def seedvr2_7b() -> "ModelConfig":
+        return AVAILABLE_MODELS["seedvr2-7b"]
+
     def x_embedder_input_dim(self) -> int:
         if "Fill" in self.model_name:
             return 384
@@ -468,4 +473,28 @@ def from_name(
         supports_guidance=True,
         requires_sigma_shift=None,
     ),
+    "seedvr2-7b": ModelConfig(
+        priority=21,
+        aliases=["seedvr2-7b", "seedvr2-7B"],
+        model_name="numz/SeedVR2_comfyUI",
+        base_model=None,
+        controlnet_model=None,
+        custom_transformer_model=None,
+        num_train_steps=None,
+        max_sequence_length=None,
+        supports_guidance=True,
+        requires_sigma_shift=None,
+        transformer_overrides={
+            "vid_dim": 3072,
+            "heads": 24,
+            "num_layers": 36,
+            "mm_layers": 36,
+            "rope_dim": 64,
+            "rope_on_text": False,
+            "rope_freqs_for": "pixel",
+            "mlp_type": "normal",
+            "use_output_ada": False,
+            "last_layer_vid_only": False,
+        },
+    ),
 }

src/mflux/models/seedvr2/README.md

@@ -11,6 +11,7 @@ SeedVR2 is more recent and the preferred method for high-fidelity upscaling and
 
 ```sh
 mflux-upscale-seedvr2 \
+  --model seedvr2-7b \
   --image-path "input.png" \
   --resolution 2160 \
   --softness 0.5
@@ -34,7 +35,7 @@ image.save("input_upscaled.png")
 ```
 </details>
 
-This will upscale the image such that the shortest side is 2160 pixels while maintaining the aspect ratio.
+This will upscale the image such that the shortest side is 2160 pixels while maintaining the aspect ratio. If `--model` is omitted, MFLUX defaults to `seedvr2-3b`.
 
 Instead of specifying a target resolution, you can also use `--resolution 2x` or `--resolution 3x` to upscale by a factor of 2 or 3 respectively.
 
@@ -145,4 +146,3 @@ image.save("image_upscaled.png")
 </details>
 
 </details>
-

src/mflux/models/seedvr2/cli/seedvr2_upscale.py

@@ -23,6 +23,32 @@ def _is_image_file(path: Path) -> bool:
     return path.is_file() and path.suffix.lower() in SUPPORTED_IMAGE_SUFFIXES
 
 
+def _resolve_seedvr2_model(model_arg: str | None, model_path: str | None) -> tuple[ModelConfig, str | None]:
+    if model_arg is None:
+        return ModelConfig.seedvr2_3b(), model_path
+
+    normalized = model_arg.lower()
+    if normalized in {"seedvr2", "seedvr2-3b"}:
+        return ModelConfig.seedvr2_3b(), None
+    if normalized in {"seedvr2-7b"}:
+        return ModelConfig.seedvr2_7b(), None
+
+    if model_path is not None:
+        path = Path(model_path).expanduser()
+        if path.is_dir():
+            has_3b = (path / "seedvr2_ema_3b_fp16.safetensors").exists()
+            has_7b = (path / "seedvr2_ema_7b_fp16.safetensors").exists()
+            if has_7b and not has_3b:
+                return ModelConfig.seedvr2_7b(), model_path
+            if has_3b and not has_7b:
+                return ModelConfig.seedvr2_3b(), model_path
+
+    source = (model_path or model_arg).lower()
+    if "seedvr2_ema_7b" in source or "seedvr2-7b" in source:
+        return ModelConfig.seedvr2_7b(), model_path
+    return ModelConfig.seedvr2_3b(), model_path
+
+
 def _expand_image_paths(image_paths: list[Path]) -> list[Path]:
     expanded: list[Path] = []
     for image_path in image_paths:
@@ -53,11 +79,13 @@ def main():
         print("No images to upscale.")
         return
 
+    model_config, resolved_model_path = _resolve_seedvr2_model(args.model, args.model_path)
+
     # 3. Load the SeedVR2 model
     model = SeedVR2(
         quantize=args.quantize,
-        model_path=args.model_path,
-        model_config=ModelConfig.seedvr2_3b(),
+        model_path=resolved_model_path,
+        model_config=model_config,
     )
 
     # 4. Register callbacks

src/mflux/models/seedvr2/model/seedvr2_transformer/attention.py

@@ -16,6 +16,8 @@ def __init__(
         qk_bias: bool = False,
         qk_norm_eps: float = 1e-5,
         rope_dim: int = 128,
+        rope_freqs_for: str = "lang",
+        rope_on_text: bool = True,
         shared_weights: bool = False,
         window: tuple[int, int, int] = (4, 3, 3),
         shift: bool = False,
@@ -27,6 +29,7 @@ def __init__(
         self.scale = head_dim**-0.5
         self.window = window
         self.shift = shift
+        self.rope_on_text = rope_on_text
 
         inner_dim = heads * head_dim
 
@@ -46,7 +49,7 @@ def __init__(
             self.norm_q_txt = RMSNorm(head_dim, eps=qk_norm_eps)
             self.norm_k_txt = RMSNorm(head_dim, eps=qk_norm_eps)
 
-        self.rope = RoPEModule(dim=rope_dim)
+        self.rope = RoPEModule(dim=rope_dim, freqs_for=rope_freqs_for)
 
     def __call__(self, vid, txt, vid_shape, txt_shape):
         B, L, Bt, Lt = vid.shape[0], vid.shape[1], txt.shape[0], txt.shape[1]
@@ -67,14 +70,21 @@ def __call__(self, vid, txt, vid_shape, txt_shape):
         q_txt_rep, k_txt_rep, v_txt_rep = qkv_t_rep[:, 0], qkv_t_rep[:, 1], qkv_t_rep[:, 2]
 
         # 3. Apply RoPE
-        q_vid, k_vid, q_txt_rep, k_txt_rep = self.rope(
-            vid_q=q_vid,
-            vid_k=k_vid,
-            vid_shape=partitioner.window_shapes,
-            txt_q=q_txt_rep,
-            txt_k=k_txt_rep,
-            txt_shape=mx.repeat(txt_shape, mx.array(counts), axis=0),
-        )
+        if self.rope_on_text:
+            q_vid, k_vid, q_txt_rep, k_txt_rep = self.rope(
+                vid_q=q_vid,
+                vid_k=k_vid,
+                vid_shape=partitioner.window_shapes,
+                txt_q=q_txt_rep,
+                txt_k=k_txt_rep,
+                txt_shape=mx.repeat(txt_shape, mx.array(counts), axis=0),
+            )
+        else:
+            q_vid, k_vid = self.rope(
+                vid_q=q_vid,
+                vid_k=k_vid,
+                vid_shape=partitioner.window_shapes,
+            )
 
         # 4. Attention
         vid_lens = mx.prod(partitioner.window_shapes, axis=1)

src/mflux/models/seedvr2/model/seedvr2_transformer/mm_swiglu.py

@@ -1,7 +1,7 @@
 import mlx.core as mx
 from mlx import nn
 
-from mflux.models.seedvr2.model.seedvr2_transformer.swiglu_mlp import SwiGLUMLP
+from mflux.models.seedvr2.model.seedvr2_transformer.swiglu_mlp import GELUMLP, SwiGLUMLP
 
 
 class MMSwiGLU(nn.Module):
@@ -12,17 +12,25 @@ def __init__(
         expand_ratio: int = 4,
         shared_weights: bool = False,
         is_last_layer: bool = False,
+        mlp_type: str = "swiglu",
     ):
         super().__init__()
         self.shared_weights = shared_weights
         self.is_last_layer = is_last_layer
+        self.mlp_type = mlp_type
+
+        mlp_cls = SwiGLUMLP
+        mlp_kwargs = {"expand_ratio": expand_ratio}
+        if mlp_type == "normal":
+            mlp_cls = GELUMLP
+            mlp_kwargs["bias"] = True
 
         if shared_weights:
-            self.all = SwiGLUMLP(dim=vid_dim, expand_ratio=expand_ratio)
+            self.all = mlp_cls(dim=vid_dim, **mlp_kwargs)
         else:
-            self.vid = SwiGLUMLP(dim=vid_dim, expand_ratio=expand_ratio)
+            self.vid = mlp_cls(dim=vid_dim, **mlp_kwargs)
             if not is_last_layer:
-                self.txt = SwiGLUMLP(dim=txt_dim, expand_ratio=expand_ratio)
+                self.txt = mlp_cls(dim=txt_dim, **mlp_kwargs)
 
     def __call__(self, vid: mx.array, txt: mx.array) -> tuple[mx.array, mx.array]:
         if self.shared_weights: