tiny-anima-modular-pipe

Tiny randomly-initialized Anima modular pipeline, used by the diffusers fast tests in tests/modular_pipelines/anima/. Not useful for generation.

tokenizer and t5_tokenizer are not stored here — modular_model_index.json points them at hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration and hf-internal-testing/tiny-random-t5.

Build script

"""Build the tiny Anima modular fixture repository."""

import json
import os
import shutil
import sys

import torch
from transformers import Qwen2Tokenizer, Qwen3Config, Qwen3Model, T5TokenizerFast

from diffusers import (
    AnimaAutoBlocks,
    AnimaTextConditioner,
    AutoencoderKLQwenImage,
    CosmosTransformer3DModel,
    FlowMatchEulerDiscreteScheduler,
)


REPO_ID = "hf-internal-testing/tiny-anima-modular-pipe"
OUT = sys.argv[1]


def get_dummy_components():
    torch.manual_seed(0)
    transformer = CosmosTransformer3DModel(
        in_channels=4,
        out_channels=4,
        num_attention_heads=2,
        attention_head_dim=16,
        num_layers=2,
        mlp_ratio=2,
        text_embed_dim=16,
        adaln_lora_dim=4,
        max_size=(4, 32, 32),
        patch_size=(1, 2, 2),
        rope_scale=(1.0, 4.0, 4.0),
        concat_padding_mask=True,
        extra_pos_embed_type=None,
    )

    torch.manual_seed(0)
    vae = AutoencoderKLQwenImage(
        base_dim=24,
        z_dim=4,
        dim_mult=[1, 2, 4],
        num_res_blocks=1,
        temperal_downsample=[False, True],
        latents_mean=[0.0] * 4,
        latents_std=[1.0] * 4,
    )

    torch.manual_seed(0)
    text_conditioner = AnimaTextConditioner(
        source_dim=16,
        target_dim=16,
        model_dim=16,
        num_layers=2,
        num_attention_heads=4,
        target_vocab_size=32128,
        min_sequence_length=16,
    )

    torch.manual_seed(0)
    text_encoder_config = Qwen3Config(
        vocab_size=152064,
        hidden_size=16,
        intermediate_size=32,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=2,
        max_position_embeddings=128,
        rms_norm_eps=1e-6,
        rope_theta=1000000.0,
        head_dim=4,
        attention_bias=False,
    )
    text_encoder = Qwen3Model(text_encoder_config).eval()
    tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration")
    t5_tokenizer = T5TokenizerFast.from_pretrained("hf-internal-testing/tiny-random-t5")
    scheduler = FlowMatchEulerDiscreteScheduler(shift=3.0)

    return {
        "transformer": transformer,
        "vae": vae,
        "scheduler": scheduler,
        "text_encoder": text_encoder,
        "tokenizer": tokenizer,
        "t5_tokenizer": t5_tokenizer,
        "text_conditioner": text_conditioner,
    }


if os.path.isdir(OUT):
    shutil.rmtree(OUT)

pipe = AnimaAutoBlocks().init_pipeline()
pipe.update_components(**get_dummy_components())
pipe.save_pretrained(OUT, safe_serialization=True)

index_path = os.path.join(OUT, "modular_model_index.json")
with open(index_path) as f:
    index = json.load(f)

# the tokenizers are not retrained, so the fixture points at the repositories they come from instead of
# duplicating their files. Their class names are the ones the Anima blocks declare, which resolve on both
# transformers 4.x (where `T5TokenizerFast` is the fast class) and 5.x (where it aliases `T5Tokenizer`).
TOKENIZER_SOURCES = {
    "tokenizer": ("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration", "Qwen2Tokenizer"),
    "t5_tokenizer": ("hf-internal-testing/tiny-random-t5", "T5TokenizerFast"),
}

for name, entry in index.items():
    if not isinstance(entry, list):
        continue
    spec = entry[2]
    assert spec["pretrained_model_name_or_path"] == OUT, (name, spec)
    if name in TOKENIZER_SOURCES:
        repo, class_name = TOKENIZER_SOURCES[name]
        entry[1] = class_name
        spec["pretrained_model_name_or_path"] = repo
        spec["subfolder"] = None
        spec["type_hint"] = ["transformers", class_name]
        shutil.rmtree(os.path.join(OUT, name))
    else:
        spec["pretrained_model_name_or_path"] = REPO_ID

with open(index_path, "w") as f:
    json.dump(index, f, indent=2, sort_keys=True)

print(json.dumps(index, indent=2, sort_keys=True))
print("\n".join(sorted(str(p) for p in os.listdir(OUT))))