README.md

metadata

library_name: diffusers
base_model:
  - OpenMOSS-Team/MOSS-Image-Base

This tiny model is intended for debugging. It is randomly initialized using the configuration adapted from OpenMOSS-Team/MOSS-Image-Base.

Structural comparison with OpenMOSS-Team/MOSS-Image-Base:

Feature	Source	Tiny
Pipeline	custom MOSS Image	same
Text vocabulary / precision	248,320 / BF16	same
Qwen3.5 text layers and cadence	32; 3 linear + 1 full	8; two cycles
Text width / attention head dimension	4,096 / 256	32 / 32
Text vision tower	27 layers	2 layers; width 32
DiT blocks	42	4
DiT width / heads / FFN	3,072 / 24 / 8,192	64 / 2 / 128
DiT latent channels	128	same
Flux2 VAE blocks / layers per block	4 / 2	same; widths reduced
Checkpoint precision	BF16 text/DiT; FP32 VAE/RoPE	same

The source's custom scheduler, DiT, tokenizer and pipeline code are copied from the pinned revision. The transformer keeps four repeated blocks, and the VAE stays FP32 because that is the source checkpoint representation and decoding path.

File path	Size
text_encoder/model.safetensors	16.4MB
transformer/diffusion_pytorch_model.safetensors	0.6MB
vae/diffusion_pytorch_model.safetensors	5.8MB

Example usage:

• Transformers (✅ verified)

import torch
from diffusers import DiffusionPipeline

model_id = "tiny-random/moss-image"
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = DiffusionPipeline.from_pretrained(
    model_id,
    custom_pipeline="pipeline_moss_image",
    trust_remote_code=True,
    torch_dtype={"default": torch.bfloat16, "vae": torch.float32},
)
pipe.to(device)
image = pipe(
    "A red fox in a snowy pine forest",
    height=32,
    width=32,
    num_inference_steps=1,
    guidance_scale=1.0,
    uncond_skip_layers=(2,),
    generator=torch.Generator("cpu").manual_seed(42),
).images[0]
print(image.size)

Codes to create this repo (✅ verified):

Click to expand

import importlib.util
import json
from pathlib import Path
from shutil import copyfile

import torch
from diffusers import AutoencoderKLFlux2
from huggingface_hub import hf_hub_download, list_repo_files
from transformers import AutoConfig, Qwen3_5Model, set_seed

source_model_id = "OpenMOSS-Team/MOSS-Image-Base"
save_folder = "/tmp/tiny-random/moss-image"
source_revision = "85975793d8b3d432f1b6e27f352abaa4c2a4c0d3"
set_seed(42)
folder = Path(save_folder)
folder.mkdir(parents=True, exist_ok=True)

def load_source_module():
    path = hf_hub_download(source_model_id, "pipeline_moss_image.py", revision=source_revision)
    spec = importlib.util.spec_from_file_location("moss_image_source", path)
    module = importlib.util.module_from_spec(spec)
    if spec.loader is None:
        raise RuntimeError("The pinned MOSS pipeline has no import loader")
    spec.loader.exec_module(module)
    return module

def initialize(module):
    with torch.no_grad():
        for name, parameter in sorted(module.named_parameters()):
            if not parameter.is_floating_point():
                continue
            if name.endswith("A_log") or name.endswith("dt_bias"):
                continue
            if name.startswith("language_model") and (
                name.endswith("input_layernorm.weight")
                or name.endswith("post_attention_layernorm.weight")
                or name.endswith(".linear_attn.norm.weight")
                or name.endswith("language_model.norm.weight")
                or name.endswith(".self_attn.q_norm.weight")
                or name.endswith(".self_attn.k_norm.weight")
            ):
                parameter.zero_()
            elif ("norm" in name or ".bn." in name) and name.endswith("weight"):
                parameter.fill_(1)
            else:
                torch.nn.init.normal_(parameter, mean=0.0, std=0.02)

source = load_source_module()

raw_text = json.loads(Path(hf_hub_download(
    source_model_id, "text_encoder/config.json", revision=source_revision,
)).read_text(encoding="utf-8"))
text = raw_text["text_config"]
text.update(
    hidden_size=32,
    head_dim=32,
    num_attention_heads=1,
    num_key_value_heads=1,
    intermediate_size=64,
    linear_key_head_dim=32,
    linear_value_head_dim=32,
    linear_num_key_heads=1,
    linear_num_value_heads=1,
    num_hidden_layers=8,
    layer_types=["linear_attention", "linear_attention", "linear_attention", "full_attention"] * 2,
)
text["rope_parameters"]["mrope_section"] = [1, 1, 2]
raw_text["vision_config"].update(
    hidden_size=32,
    intermediate_size=64,
    num_heads=1,
    out_hidden_size=32,
    depth=2,
)
text_folder = folder / "text_encoder"
text_folder.mkdir(exist_ok=True)
(text_folder / "config.json").write_text(json.dumps(raw_text, indent=2), encoding="utf-8")
text_encoder = Qwen3_5Model(AutoConfig.from_pretrained(text_folder)).to(torch.bfloat16)
initialize(text_encoder)
text_encoder.save_pretrained(text_folder, safe_serialization=True)

# The source stores a 42-block custom DiT in BF16.  Its VAE-facing 128
# channels remain fixed: the pipeline unshuffles them to the VAE's 32
# latent channels before decoding.
transformer = source.MOSSImageTransformer(
    dim=64,
    ffn_dim=128,
    num_heads=2,
    num_layers=4,
    text_dim=32,
    freq_dim=32,
    latent_channels=128,
).to(torch.bfloat16)
initialize(transformer)
transformer.save_pretrained(folder / "transformer", safe_serialization=True)

raw_vae = json.loads(Path(hf_hub_download(
    source_model_id, "vae/config.json", revision=source_revision,
)).read_text(encoding="utf-8"))
raw_vae.update(
    block_out_channels=[32, 32, 64, 64],
    decoder_block_out_channels=[32, 32, 64, 64],
    norm_num_groups=32,
    sample_size=32,
)
vae = AutoencoderKLFlux2.from_config(raw_vae).to(torch.float32)
initialize(vae)
vae.save_pretrained(folder / "vae", safe_serialization=True)

scheduler = source.MOSSImageScheduler()
scheduler.save_pretrained(folder / "scheduler")

# Preserve the executable source and tokenizer assets.  The source pipeline
# defaults its negative-branch skip to block 22 of a 42-block DiT.  Map that
# central source block to block 2 of the retained four-block topology so a
# normal public ``pipe(prompt)`` call remains valid after depth reduction.
assets = {
    "model_index.json", "pipeline_moss_image.py", "scheduler/pipeline_moss_image.py",
    "scheduler/scheduler_config.json", "tokenizer/chat_template.jinja",
    "tokenizer/tokenizer.json", "tokenizer/tokenizer_config.json",
    "transformer/pipeline_moss_image.py",
}
for filename in sorted(set(list_repo_files(source_model_id, revision=source_revision)) & assets):
    destination = folder / filename
    destination.parent.mkdir(parents=True, exist_ok=True)
    copyfile(hf_hub_download(source_model_id, filename, revision=source_revision), destination)
for pipeline_path in (
    folder / "pipeline_moss_image.py",
    folder / "scheduler" / "pipeline_moss_image.py",
    folder / "transformer" / "pipeline_moss_image.py",
):
    source_code = pipeline_path.read_text(encoding="utf-8")
    source_code = source_code.replace(
        "uncond_skip_layers=(22,)", "uncond_skip_layers=(2,)"
    ).replace(
        "defaults to (22,)", "defaults to (2,)"
    )
    pipeline_path.write_text(source_code, encoding="utf-8")

pipe = source.MOSSImagePipeline(transformer, scheduler, text_encoder, None, vae)

Printing the model:

Click to expand

MOSSImagePipeline {
  "_class_name": "MOSSImagePipeline",
  "_diffusers_version": "0.40.0.dev0",
  "scheduler": [
    "moss_image_source",
    "MOSSImageScheduler"
  ],
  "text_encoder": [
    "transformers",
    "Qwen3_5Model"
  ],
  "tokenizer": [
    null,
    null
  ],
  "transformer": [
    "moss_image_source",
    "MOSSImageTransformer"
  ],
  "vae": [
    "diffusers",
    "AutoencoderKLFlux2"
  ]
}

Test environment:

• Pillow: 12.3.0
• diffusers: 0.40.0.dev0
• huggingface_hub: 1.27.0
• safetensors: 0.8.0
• tokenizers: 0.23.0rc0
• torch: 2.13.0+cu126
• transformers: 5.15.0.dev0

Change log:

• 2026-10-23: Initial version.

This repo is generated by GPT-6.1-Sol