README.md

---
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):

<details><summary>Click to expand</summary>

```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)
```

</details>

### Printing the model:

<details><summary>Click to expand</summary>

```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"
  ]
}
```

</details>

### 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