---
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](https://huggingface.co/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)
```python
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
```python
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)
```
Click to expand
```text
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"
]
}
```