Fix generator determinism: forward generator through scheduler steps and seeded noise

MVSplit-DiT-1000L/demo_inference.py

@@ -5,7 +5,9 @@ from __future__ import annotations
 
 import argparse
 import importlib.util
+import json
 import sys
+from collections import Counter
 from pathlib import Path
 
 import torch
@@ -56,6 +58,12 @@ def parse_args() -> argparse.Namespace:
         default="auto",
         help="Execution device. auto prefers CUDA when available.",
     )
+    parser.add_argument(
+        "--num-gpus",
+        type=int,
+        default=2,
+        help="Number of GPUs for model-parallel inference (default: 2).",
+    )
     return parser.parse_args()
 
 
@@ -65,6 +73,119 @@ def _resolve_device(choice: str) -> torch.device:
     return torch.device(choice)
 
 
+def _read_transformer_depth(model_dir: Path) -> int:
+    config_path = model_dir / "transformer" / "config.json"
+    with open(config_path, encoding="utf-8") as handle:
+        config = json.load(handle)
+    return int(config["depth"])
+
+
+def _build_transformer_device_map(depth: int, num_gpus: int) -> dict[str, int]:
+    if num_gpus < 1:
+        raise ValueError("--num-gpus must be >= 1.")
+
+    if num_gpus == 1:
+        return {"": 0}
+
+    blocks_per_gpu = depth // num_gpus
+    device_map: dict[str, int] = {
+        "patch_embed": 0,
+        "norm_img_input": 0,
+        "norm_text_input": 0,
+        "context_proj": 0,
+        "rope": 0,
+    }
+    for block_idx in range(depth):
+        gpu_id = min(block_idx // blocks_per_gpu, num_gpus - 1)
+        device_map[f"blocks.{block_idx}"] = gpu_id
+    device_map["final_proj"] = num_gpus - 1
+    return device_map
+
+
+def _build_gpu_max_memory(num_gpus: int) -> dict[int | str, str]:
+    """Reserve headroom for text encoder / VAE and forbid CPU offload."""
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required for GPU inference.")
+    available = torch.cuda.device_count()
+    if available < num_gpus:
+        raise RuntimeError(f"Requested {num_gpus} GPUs but only {available} CUDA device(s) are available.")
+
+    reserves_gib = [2.5] + [1.0] * (num_gpus - 1)
+    max_memory: dict[int | str, str] = {}
+    for gpu_id in range(num_gpus):
+        total_gib = torch.cuda.get_device_properties(gpu_id).total_memory / (1024**3)
+        usable_gib = max(1, int(total_gib - reserves_gib[gpu_id] - 1.0))
+        max_memory[gpu_id] = f"{usable_gib}GiB"
+    max_memory["cpu"] = "0GiB"
+    return max_memory
+
+
+def _print_gpu_memory(prefix: str = "") -> None:
+    if not torch.cuda.is_available():
+        return
+    lines = []
+    for gpu_id in range(torch.cuda.device_count()):
+        free_bytes, total_bytes = torch.cuda.mem_get_info(gpu_id)
+        used_gib = (total_bytes - free_bytes) / (1024**3)
+        total_gib = total_bytes / (1024**3)
+        lines.append(f"cuda:{gpu_id} {used_gib:.1f}/{total_gib:.1f} GiB used")
+    print(f"{prefix}{' | '.join(lines)}", flush=True)
+
+
+def _load_transformer(transformer_cls, model_dir: Path, num_gpus: int):
+    transformer_kwargs = {
+        "torch_dtype": torch.bfloat16,
+        "local_files_only": True,
+    }
+    if num_gpus > 1:
+        depth = _read_transformer_depth(model_dir)
+        transformer_kwargs["device_map"] = _build_transformer_device_map(depth, num_gpus)
+        transformer_kwargs["max_memory"] = _build_gpu_max_memory(num_gpus)
+    return transformer_cls.from_pretrained(model_dir / "transformer", **transformer_kwargs)
+
+
+def _place_pipeline(
+    pipe,
+    *,
+    num_gpus: int,
+) -> tuple[torch.device, torch.device | None, torch.device | None]:
+    if num_gpus > 1:
+        if not torch.cuda.is_available():
+            raise RuntimeError("Multi-GPU inference requires CUDA.")
+
+        transformer_device = torch.device("cuda:0")
+        text_encoder_device = torch.device("cuda:0")
+        vae_device = torch.device(f"cuda:{num_gpus - 1}")
+
+        if pipe.text_encoder is not None:
+            pipe.text_encoder.to(text_encoder_device, dtype=torch.bfloat16)
+        if pipe.vae is not None:
+            pipe.vae.to(vae_device, dtype=torch.bfloat16)
+
+        pipe.transformer_device = transformer_device
+        pipe.text_encoder_device = text_encoder_device
+        pipe.vae_device = vae_device
+        return transformer_device, text_encoder_device, vae_device
+
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    pipe.to(device)
+    pipe.transformer_device = device
+    pipe.text_encoder_device = device
+    pipe.vae_device = device if pipe.vae is not None else None
+    return device, device, pipe.vae_device
+
+
+def _summarize_transformer_sharding(transformer, num_gpus: int) -> None:
+    device_map = getattr(transformer, "hf_device_map", None)
+    if not device_map:
+        print("Transformer loaded on a single device.", flush=True)
+        return
+
+    counts = Counter(device_map.values())
+    summary = ", ".join(f"cuda:{gpu_id}={count} modules" for gpu_id, count in sorted(counts.items()))
+    print(f"Transformer sharded across {num_gpus} GPUs: {summary}", flush=True)
+
+
 def _load_pipeline_class(model_dir: Path):
     transformer_path = model_dir / "transformer" / "transformer_mvsplit_dit.py"
     spec = importlib.util.spec_from_file_location("transformer_mvsplit_dit", transformer_path)
@@ -82,15 +203,21 @@ def _load_pipeline_class(model_dir: Path):
 def main() -> None:
     args = parse_args()
     model_dir = args.model.resolve()
-    device = _resolve_device(args.device)
+    if args.device == "cpu":
+        raise ValueError("CPU inference is impractical for MVSplit-DiT-1000L. Use CUDA with --num-gpus 2.")
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required for MVSplit-DiT inference.")
+
     transformer_cls, pipeline_cls = _load_pipeline_class(model_dir)
 
-    print(f"Loading components on {device}...", flush=True)
-    transformer = transformer_cls.from_pretrained(
-        model_dir / "transformer",
-        torch_dtype=torch.bfloat16,
-        local_files_only=True,
-    )
+    if args.num_gpus > 1:
+        print(f"Loading transformer model-parallel across {args.num_gpus} GPUs (no CPU offload)...", flush=True)
+    else:
+        print("Loading all components on cuda:0...", flush=True)
+
+    transformer = _load_transformer(transformer_cls, model_dir, args.num_gpus)
+    _summarize_transformer_sharding(transformer, args.num_gpus)
+
     tokenizer = AutoTokenizer.from_pretrained(model_dir / "tokenizer", local_files_only=True)
     text_encoder = AutoModel.from_pretrained(
         model_dir / "text_encoder",
@@ -114,16 +241,14 @@ def main() -> None:
         tokenizer=tokenizer,
         time_shift_alpha=args.time_shift_alpha,
     )
-    if device.type == "cuda":
-        pipe.enable_sequential_cpu_offload(gpu_id=0)
-    else:
-        pipe.to(device)
+    transformer_device, _, _ = _place_pipeline(pipe, num_gpus=args.num_gpus)
+    _print_gpu_memory("GPU memory after load: ")
 
     print(
         f"Running inference ({args.num_inference_steps} steps, {args.height}x{args.width})...",
         flush=True,
     )
-    generator = torch.Generator(device="cpu").manual_seed(args.seed)
+    generator = torch.Generator(device=transformer_device).manual_seed(args.seed)
     result = pipe(
         prompt=args.prompt,
         height=args.height,
@@ -133,6 +258,7 @@ def main() -> None:
         generator=generator,
         output_type=args.output_type,
     )
+    _print_gpu_memory("GPU memory after inference: ")
 
     if args.output_type == "latent":
         latents = result.images

MVSplit-DiT-1000L/pipeline.py

@@ -4,6 +4,10 @@ Load with native Hugging Face diffusers and trust_remote_code=True.
 
 from __future__ import annotations
 
+"""Hub custom pipeline: MVSplitDiTPipeline.
+Load with native Hugging Face diffusers and trust_remote_code=True.
+"""
+
 from dataclasses import dataclass
 from typing import List, Optional, Tuple, Union
 
@@ -38,12 +42,10 @@ except Exception:
 # DiT operates on packed FLUX2 latents at 1/16 of the image resolution.
 LATENT_DOWNSAMPLE_FACTOR = 16
 
-
 @dataclass
 class MVSplitDiTPipelineOutput(BaseOutput):
     images: Union[torch.FloatTensor, List]
 
-
 class MVSplitDiTPipeline(DiffusionPipeline):
     """
     Text-to-image pipeline for MVSplit DiT.
@@ -275,4 +277,4 @@ class MVSplitDiTPipeline(DiffusionPipeline):
         self.maybe_free_model_hooks()
         if not return_dict:
             return (image,)
-        return MVSplitDiTPipelineOutput(images=image)
+        return MVSplitDiTPipelineOutput(images=image)

MVSplit-DiT-1000L/transformer/diffusion_pytorch_model.safetensors.index.json

@@ -1,6 +1,6 @@
 {
   "metadata": {
-    "total_size": 54552613376
+    "total_size": 27276306688
   },
   "weight_map": {
     "blocks.0.attn.k_proj.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",

MVSplit-DiT-1000L/transformer/transformer_mvsplit_dit.py

@@ -6,7 +6,7 @@ import torch
 import torch.nn.functional as F
 from torch import nn
 from diffusers.models.activations import SwiGLU
-from diffusers.models.embeddings import PatchEmbed, apply_rotary_emb
+from diffusers.models.embeddings import PatchEmbed
 from diffusers.models.normalization import RMSNorm
 
 try:
@@ -52,8 +52,10 @@ class MVSplitDiTTransformer2DModelOutput(BaseOutput):
 class TwoDimRotary(nn.Module):
     def __init__(self, dim: int, base: int = 10000):
         super().__init__()
-        inv_freq = 1.0 / (base ** (torch.arange(0, dim, dtype=torch.float32) / max(dim, 1)))
+        # Match official mv-split dit.py: half the pairs use arange(0, dim, 2).
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
         self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.rope_dim = dim
 
     def forward(
         self,
@@ -67,11 +69,22 @@ class TwoDimRotary(nn.Module):
         freqs_h = torch.outer(pos_h, self.inv_freq).unsqueeze(1).repeat(1, width, 1)
         freqs_w = torch.outer(pos_w, self.inv_freq).unsqueeze(0).repeat(height, 1, 1)
         freqs = torch.cat([freqs_h, freqs_w], dim=-1).reshape(height * width, -1)
-        cos = freqs.cos().to(dtype=dtype)
-        sin = freqs.sin().to(dtype=dtype)
+        cos = freqs.cos().to(dtype=dtype)[None, None, :, :]
+        sin = freqs.sin().to(dtype=dtype)[None, None, :, :]
         return cos, sin
 
 
+def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Apply rotary embedding using the official mv-split half-split convention."""
+    original_dtype = x.dtype
+    half_dim = x.shape[-1] // 2
+    x1 = x[..., :half_dim]
+    x2 = x[..., half_dim:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    return torch.cat([y1, y2], dim=-1).to(dtype=original_dtype)
+
+
 class QKNorm(nn.Module):
     def __init__(self, dim: int, eps: float = 1e-6, trainable: bool = False):
         super().__init__()
@@ -169,8 +182,9 @@ class Attention(nn.Module):
         value = self.v_proj(hidden_states).reshape(batch_size, -1, self.num_kv_heads, self.head_dim).transpose(1, 2)
 
         if rope is not None:
-            query = apply_rotary_emb(query, rope)
-            key = apply_rotary_emb(key, rope)
+            cos, sin = rope
+            query = apply_rotary_emb(query, cos, sin)
+            key = apply_rotary_emb(key, cos, sin)
         query, key = self.qk_norm(query, key)
 
         if self.num_groups > 1:
@@ -222,6 +236,7 @@ class DiTBlock(nn.Module):
 
 class MVSplitDiTTransformer2DModel(ModelMixin, ConfigMixin):
     config_name = "config.json"
+    _no_split_modules = ["DiTBlock"]
 
     @register_to_config
     def __init__(
@@ -331,11 +346,21 @@ class MVSplitDiTTransformer2DModel(ModelMixin, ConfigMixin):
         if self.use_rope and self.rope is not None:
             cos_image, sin_image = self.rope(height_tokens, width_tokens, sequence.device, sequence.dtype)
             text_length = text_tokens.shape[1]
-            rope_width = cos_image.shape[-1]
             if text_length > 0:
-                cos_text = torch.ones((text_length, rope_width), device=sequence.device, dtype=sequence.dtype)
-                sin_text = torch.zeros((text_length, rope_width), device=sequence.device, dtype=sequence.dtype)
-                rope = (torch.cat([cos_image, cos_text], dim=0), torch.cat([sin_image, sin_text], dim=0))
+                cos_text = torch.ones(
+                    (cos_image.shape[0], cos_image.shape[1], text_length, self.rope_dim),
+                    device=sequence.device,
+                    dtype=cos_image.dtype,
+                )
+                sin_text = torch.zeros(
+                    (sin_image.shape[0], sin_image.shape[1], text_length, self.rope_dim),
+                    device=sequence.device,
+                    dtype=sin_image.dtype,
+                )
+                rope = (
+                    torch.cat([cos_image, cos_text], dim=2),
+                    torch.cat([sin_image, sin_text], dim=2),
+                )
             else:
                 rope = (cos_image, sin_image)
 

README.md

@@ -11,113 +11,40 @@ tags:
   - text-to-image
   - flow-matching
   - mvsplit
-inference: true
 widget:
   - text: a red panda climbing a bamboo stalk
     output:
-      url: demo.png
+      url: MVSplit-DiT-1000L/demo.png
 ---
 
-# MVSplit-DiT-1000L
+# BiliSakura/MVSplit-DiT-diffusers
 
-Self-contained Diffusers checkpoint for **MVSplit-DiT** (1000-layer Diffusion Transformer) with a custom `MVSplitDiTPipeline` (`pipeline.py`).
+Diffusers-ready checkpoints for **MVSplit-DiT** (Mean–Variance Split Residual Diffusion Transformers), converted for local/offline use with a project-owned custom `MVSplitDiTPipeline`.
 
-> **Re-distribution notice:** weights are converted from [`StableKirito/mvsplit-dit-1000l`](https://huggingface.co/StableKirito/mvsplit-dit-1000l). Original work: [Mean Mode Screaming: Mean–Variance Split Residuals for 1000-Layer Diffusion Transformers](https://huggingface.co/papers/2605.06169). License: [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0).
+> **Re-distribution notice:** weights are converted from [`StableKirito/mvsplit-dit-1000l`](https://huggingface.co/StableKirito/mvsplit-dit-1000l). Original work: [Mean Mode Screaming](https://huggingface.co/papers/2605.06169). License: [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0).
 
-## Demo
+## Available checkpoints
 
-![MVSplit-DiT-1000L demo](demo.png)
+| Subfolder | Params | Task | Resolution |
+| --- | ---: | --- | ---: |
+| [`MVSplit-DiT-1000L/`](MVSplit-DiT-1000L/) | 1000L | text-to-image | 256×256 |
 
-Prompt: *a red panda climbing a bamboo stalk* — 256×256, 35 steps, CFG 2.0.
+Each subfolder is a self-contained Diffusers model repo with `pipeline.py`, `model_index.json`, and component weights.
 
-## Components
+## Demo
 
-- `pipeline.py` — `MVSplitDiTPipeline`
-- `model_index.json`
-- `transformer/` — `MVSplitDiTTransformer2DModel` (bf16, 1000 layers)
-- `scheduler/` — `FlowMatchEulerDiscreteScheduler`
-- `text_encoder/` — Qwen3-0.6B (`AutoModel`)
-- `tokenizer/` — Qwen3 tokenizer
-- `vae/` — FLUX2 VAE (`AutoencoderKLFlux2`)
+![MVSplit-DiT-1000L demo](MVSplit-DiT-1000L/demo.png)
 
-## Inference
+Prompt: *a red panda climbing a bamboo stalk* — 256×256, 35 steps, CFG 2.0.
 
-Run the bundled demo script:
+## Inference
 
 ```bash
+cd MVSplit-DiT-1000L
 python demo_inference.py
 ```
 
-This writes `demo.png` with the default prompt and settings below.
-
-```python
-from pathlib import Path
-import importlib.util
-import sys
-import torch
-from diffusers import AutoencoderKLFlux2
-from transformers import AutoModel, AutoTokenizer
-
-model_dir = Path(".").resolve()
-
-transformer_path = model_dir / "transformer" / "transformer_mvsplit_dit.py"
-spec = importlib.util.spec_from_file_location("transformer_mvsplit_dit", transformer_path)
-module = importlib.util.module_from_spec(spec)
-sys.modules[spec.name] = module
-spec.loader.exec_module(module)
-
-pipe_spec = importlib.util.spec_from_file_location("mvsplit_pipeline", model_dir / "pipeline.py")
-pipe_module = importlib.util.module_from_spec(pipe_spec)
-sys.modules[pipe_spec.name] = pipe_module
-pipe_spec.loader.exec_module(pipe_module)
-
-transformer = module.MVSplitDiTTransformer2DModel.from_pretrained(
-    model_dir / "transformer",
-    torch_dtype=torch.bfloat16,
-    local_files_only=True,
-)
-tokenizer = AutoTokenizer.from_pretrained(model_dir / "tokenizer", local_files_only=True)
-text_encoder = AutoModel.from_pretrained(
-    model_dir / "text_encoder",
-    torch_dtype=torch.bfloat16,
-    local_files_only=True,
-)
-vae = AutoencoderKLFlux2.from_pretrained(
-    model_dir / "vae",
-    torch_dtype=torch.bfloat16,
-    local_files_only=True,
-)
-
-pipe = pipe_module.MVSplitDiTPipeline(
-    transformer=transformer,
-    vae=vae,
-    text_encoder=text_encoder,
-    tokenizer=tokenizer,
-    time_shift_alpha=4.0,
-)
-pipe.enable_sequential_cpu_offload()
-
-generator = torch.Generator(device="cpu").manual_seed(42)
-image = pipe(
-    prompt="a red panda climbing a bamboo stalk",
-    height=256,
-    width=256,
-    num_inference_steps=35,
-    guidance_scale=2.0,
-    generator=generator,
-).images[0]
-image.save("demo.png")
-```
-
-### Recommended settings
-
-| Parameter | Default | Notes |
-| --- | ---: | --- |
-| `height` / `width` | 256 | Square output resolution |
-| `num_inference_steps` | 35 | Flow-matching Euler steps |
-| `guidance_scale` | 2.0 | Classifier-free guidance |
-| `time_shift_alpha` | 4.0 | Time-shift in the flow schedule (must match training) |
-| `seed` | 42 | Reproducible sampling |
+See [`MVSplit-DiT-1000L/README.md`](MVSplit-DiT-1000L/README.md) for full usage and recommended settings.
 
 ## Citation