Upload folder using huggingface_hub

README.md

@@ -12,23 +12,23 @@ Structural comparison with `OpenMOSS-Team/MOSS-Image-Base`:
 |---|---:|---:|
 | Pipeline | custom MOSS Image | same |
 | Text vocabulary / precision | 248,320 / BF16 | same |
-| Qwen3.5 text layers and cadence | 32; 3 linear + 1 full | 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 | same depth; width 32 |
-| DiT blocks | 42 | same |
+| 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 all 42 blocks, and the VAE stays
-FP32 because that is the source checkpoint representation and decoding path.
+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 | 17.5MB |
-| transformer/diffusion_pytorch_model.safetensors | 3.8MB |
+| text_encoder/model.safetensors | 16.4MB |
+| transformer/diffusion_pytorch_model.safetensors | 0.6MB |
 | vae/diffusion_pytorch_model.safetensors | 5.8MB |
 
 
@@ -55,6 +55,7 @@ image = pipe(
     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)
@@ -128,8 +129,8 @@ text.update(
     linear_value_head_dim=32,
     linear_num_key_heads=1,
     linear_num_value_heads=1,
-    num_hidden_layers=32,
-    layer_types=["linear_attention", "linear_attention", "linear_attention", "full_attention"] * 8,
+    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(
@@ -137,7 +138,7 @@ raw_text["vision_config"].update(
     intermediate_size=64,
     num_heads=1,
     out_hidden_size=32,
-    depth=27,
+    depth=2,
 )
 text_folder = folder / "text_encoder"
 text_folder.mkdir(exist_ok=True)
@@ -153,7 +154,7 @@ transformer = source.MOSSImageTransformer(
     dim=64,
     ffn_dim=128,
     num_heads=2,
-    num_layers=42,
+    num_layers=4,
     text_dim=32,
     freq_dim=32,
     latent_channels=128,
@@ -177,7 +178,10 @@ vae.save_pretrained(folder / "vae", safe_serialization=True)
 scheduler = source.MOSSImageScheduler()
 scheduler.save_pretrained(folder / "scheduler")
 
-# Preserve all executable source code plus tokenizer assets byte-for-byte.
+# 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",
@@ -188,6 +192,18 @@ for filename in sorted(set(list_repo_files(source_model_id, revision=source_revi
     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)
 ```

pipeline_moss_image.py

@@ -341,11 +341,11 @@ class MOSSImagePipeline(DiffusionPipeline):
                  generator=None, latents=None, prompt_embeds=None, prompt_attention_mask=None,
                  negative_prompt_embeds=None, negative_prompt_attention_mask=None,
                  max_sequence_length=4096, output_type='pil', return_dict=True,
-                 uncond_skip_layers=(22,), slg_layers=None, slg_scale=0.0,
+                 uncond_skip_layers=(2,), slg_layers=None, slg_scale=0.0,
                  slg_start=0.0, slg_stop=1.0):
         """Generate images with optional full-block skipping (zero-based indices).
 
-        uncond_skip_layers defaults to (22,) and replaces CFG's negative-prompt
+        uncond_skip_layers defaults to (2,) and replaces CFG's negative-prompt
         prediction with a skipped-block prediction, retaining two forwards per
         step. Pass None or an empty sequence to disable skipping. slg_layers
         adds a skipped positive-prompt forward and the fp32 correction

scheduler/pipeline_moss_image.py

@@ -341,11 +341,11 @@ class MOSSImagePipeline(DiffusionPipeline):
                  generator=None, latents=None, prompt_embeds=None, prompt_attention_mask=None,
                  negative_prompt_embeds=None, negative_prompt_attention_mask=None,
                  max_sequence_length=4096, output_type='pil', return_dict=True,
-                 uncond_skip_layers=(22,), slg_layers=None, slg_scale=0.0,
+                 uncond_skip_layers=(2,), slg_layers=None, slg_scale=0.0,
                  slg_start=0.0, slg_stop=1.0):
         """Generate images with optional full-block skipping (zero-based indices).
 
-        uncond_skip_layers defaults to (22,) and replaces CFG's negative-prompt
+        uncond_skip_layers defaults to (2,) and replaces CFG's negative-prompt
         prediction with a skipped-block prediction, retaining two forwards per
         step. Pass None or an empty sequence to disable skipping. slg_layers
         adds a skipped positive-prompt forward and the fp32 correction

text_encoder/config.json

@@ -19,30 +19,6 @@
     "initializer_range": 0.02,
     "intermediate_size": 64,
     "layer_types": [
-      "linear_attention",
-      "linear_attention",
-      "linear_attention",
-      "full_attention",
-      "linear_attention",
-      "linear_attention",
-      "linear_attention",
-      "full_attention",
-      "linear_attention",
-      "linear_attention",
-      "linear_attention",
-      "full_attention",
-      "linear_attention",
-      "linear_attention",
-      "linear_attention",
-      "full_attention",
-      "linear_attention",
-      "linear_attention",
-      "linear_attention",
-      "full_attention",
-      "linear_attention",
-      "linear_attention",
-      "linear_attention",
-      "full_attention",
       "linear_attention",
       "linear_attention",
       "linear_attention",
@@ -64,7 +40,7 @@
     "mtp_num_hidden_layers": 1,
     "mtp_use_dedicated_embeddings": false,
     "num_attention_heads": 1,
-    "num_hidden_layers": 32,
+    "num_hidden_layers": 8,
     "num_key_value_heads": 1,
     "pad_token_id": null,
     "partial_rotary_factor": 0.25,
@@ -89,7 +65,7 @@
   "video_token_id": 248057,
   "vision_config": {
     "deepstack_visual_indexes": [],
-    "depth": 27,
+    "depth": 2,
     "dtype": "bfloat16",
     "hidden_act": "gelu_pytorch_tanh",
     "hidden_size": 32,

text_encoder/model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e978b54dce1147620893ccfc0e6b4ea09e1612b318a7c7eb8c2a22fa6396495b
-size 17474184
+oid sha256:0d95a59e225028971ec117dc08ee36a90c7de8ab879ce452716abf56f225263b
+size 16417224

transformer/config.json

@@ -6,7 +6,7 @@
   "freq_dim": 32,
   "latent_channels": 128,
   "num_heads": 2,
-  "num_layers": 42,
+  "num_layers": 4,
   "rope_base": 10000.0,
   "text_dim": 32
 }

transformer/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:86214daf0570af1eeee4e1fdcba9cfe63181e538f685e652dd02663e74e5c0c7
-size 3782792
+oid sha256:ec23d054dfdc432e442046218737834d61ce0d4b33c1ad6736c8abdfb68a22e6
+size 620224

transformer/pipeline_moss_image.py

@@ -341,11 +341,11 @@ class MOSSImagePipeline(DiffusionPipeline):
                  generator=None, latents=None, prompt_embeds=None, prompt_attention_mask=None,
                  negative_prompt_embeds=None, negative_prompt_attention_mask=None,
                  max_sequence_length=4096, output_type='pil', return_dict=True,
-                 uncond_skip_layers=(22,), slg_layers=None, slg_scale=0.0,
+                 uncond_skip_layers=(2,), slg_layers=None, slg_scale=0.0,
                  slg_start=0.0, slg_stop=1.0):
         """Generate images with optional full-block skipping (zero-based indices).
 
-        uncond_skip_layers defaults to (22,) and replaces CFG's negative-prompt
+        uncond_skip_layers defaults to (2,) and replaces CFG's negative-prompt
         prediction with a skipped-block prediction, retaining two forwards per
         step. Pass None or an empty sequence to disable skipping. slg_layers
         adds a skipped positive-prompt forward and the fp32 correction

vae/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d4f21cadbc038798d0c1129ebc430f61b40bfa354f29c7d9785cefdd4b9729fa
+oid sha256:5a40fc1c545d1c3e3084b10ed0a3e9177c18508ea1d8cea3352a760079426b7b
 size 5841172