diffusion_auto.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import torch
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
from safetensors.torch import load_file

from .nextdit_crossattn import NextDiTCrossAttn, NextDiTCrossAttnConfig


@dataclass
class DiffusionConfig:
    weights_path: Optional[str] = None
    scheduler_path: Optional[str] = None
    dim: int = 1792
    n_layers: int = 24
    n_heads: int = 28
    n_kv_heads: int = 28
    latent_embedding_size: int = 3584
    input_size: int = 8
    patch_size: int = 1
    in_channels: int = 1792


@dataclass
class DiffusionBundle:
    dit: NextDiTCrossAttn
    scheduler: FlowMatchEulerDiscreteScheduler
    load_info: dict


class AutoDiffusionModel:
    """Utility to materialize a NextDiT cross-attention model and scheduler."""

    @classmethod
    def from_config(
        cls,
        config: DiffusionConfig,
        *,
        device: Optional[str] = None,
        torch_dtype: Optional[torch.dtype] = None,
    ) -> DiffusionBundle:
        dit_config = NextDiTCrossAttnConfig(
            input_size=config.input_size,
            patch_size=config.patch_size,
            in_channels=config.in_channels,
            dim=config.dim,
            n_layers=config.n_layers,
            n_heads=config.n_heads,
            n_kv_heads=config.n_kv_heads,
            latent_embedding_size=config.latent_embedding_size,
        )

        model = NextDiTCrossAttn(dit_config)
        load_info = {"missing_keys": (), "unexpected_keys": ()}

        if config.weights_path:
            weights_path = Path(config.weights_path)
            state_dict = load_file(weights_path)
            load_result = model.load_state_dict(state_dict, strict=False)
            load_info = {
                "missing_keys": load_result.missing_keys,
                "unexpected_keys": load_result.unexpected_keys,
            }

        if torch_dtype is not None:
            model = model.to(dtype=torch_dtype)
        if device is not None:
            model = model.to(device=device)

        if config.scheduler_path:
            scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(config.scheduler_path)
        else:
            scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
                "Alpha-VLLM/Lumina-Next-SFT-diffusers", subfolder="scheduler"
            )

        return DiffusionBundle(dit=model, scheduler=scheduler, load_info=load_info)


__all__ = ["DiffusionConfig", "DiffusionBundle", "AutoDiffusionModel"]