unclebane/UBIG-Zero

Uncle Bane Image Generator 2B Zero Edition (UBIG-2B-Zero)

WHAT DOES ZERO EDITION MEANS HERE!!???

It means this is a raw base model with virtually not predifined style to it, you need to bring your own data and train (full finetune) it to what you want it to look like and/or behave. So you aren't generating any Annie Mays here with this base model, only Sleep Paralysis Demons and Nightmare Fuel:

So just to drill home once again before the nerd speak starts, this model:

A modified, depth-expanded fork of NVIDIA's Cosmos-Predict2-2B-Text2Image rebuilt for low-VRAM consumer GPUs (tested on 8GB VRAM targets such as the RTX 3070).

This model replaces the monolithic legacy T5-XXL text encoder with a newer lightweight hybrid Qwen 3.5 0.8B backbone (Linear SSM / GatedDeltaNet + Gated Self-Attention) and expands the core Diffusion Transformer (DiT) from 28 layers to 40 layers using progressive depth surgery.

---

Technical Specifications

Component	Base Cosmos Predict 2 2B	This Model
DiT Depth	28 Transformer Blocks	40 Transformer Blocks (Expanded via block replication)
Text Encoder	Old T5-XXL (~4.7B parameters)	Qwen 3.5 0.8B (~800M parameters, Linear SSM hybrid)
Encoder Hidden Dim	4096	1024
Adapter Architecture	Native Cross-Attention	6-Block Cross-Attention Semantic Bridge
VRAM Baseline	> 16 GB for unquantized inference	< 7.5 GB peak VRAM (single image generation)
Precision	BFloat16 / FP8	BFloat16 / FP8 Mixed

---

Architectural Breakdown

[Prompt Text]
      │
      ▼
┌──────────────────────────────────────────────┐
│  Qwen 3.5 0.8B Text Backbone             │
│  - 24 Layers                             │
│  - Hidden Dim: 1024                      │
└──────────────────────┬───────────────────────┘
                     │ Text Embeddings (1024-dim)
                     ▼
┌──────────────────────────────────────────────┐
│  Qwen LLM Adapter (6 Transformer Blocks) │
│  - Projects 1024-dim Qwen tokens into DiT│
│  - Cross-attends with DiT sequence space │
└──────────────────────┬───────────────────────┘
                       │ Conditioned Latent Stream
                       ▼
┌──────────────────────────────────────────────┐
│ Expanded Cosmos-Predict2 DiT (40 Blocks) │
│ - Blocks 0-19: Base feature extraction   │
│ - Blocks 20-31: Cloned & shifted mid-blocks │
│ - Blocks 32-39: Re-indexed top-level blocks │
└──────────────────────────────────────────────┘

1. Depth Expansion (28 → 40 Layers)

To scale semantic capacity without retraining from scratch, the DiT backbone underwent block-expansion surgery:

• Blocks 0-19 remain anchored in their original positions to maintain base low-level and mid-level feature representations.
• Blocks 8-19 (12 mid-level semantic blocks) were duplicated and appended as Blocks 20-31 to double the model's capacity for complex compositional logic.
• Blocks 20-27 from the original 28-layer checkpoint were shifted by an offset of +12 to form the final refining Blocks 32-39.

2. Hybrid Text Encoder: Qwen 3.5 0.8B

Upstream Cosmos models rely on Old T5-XXL, which consumes 9-11 GB of VRAM on its own in FP16/BF16. By substituting Qwen3.5-0.8B:

• The entire text encoder occupies only ~1.6 GB of VRAM in BF16.
• Full-sequence context is processed using linear state-space layers (SSM GatedDeltaNet) with interleaved gated full-attention every 4 layers (blocks 3, 7, 11, 15, 19, 23).
• Prompt representation preserves fine-grained details without requiring CPU-offloading delays between sampling steps.

3. Semantic Bridge (6-Block Adapter)

Because Cosmos's native cross-attention expects conditioning aligned with diffusion steps, a 6-block LLM adapter was trained via diffusion-pipe:

• Input: 1024-dimension hidden states from Qwen 3.5 0.8B's final RMSNorm layer.
• Positional Logic: RoPE-indexed token representations projected into the DiT block streams.
• Training Regimen:
  • Stage 1: Frozen DiT backbone + Frozen Qwen 3.5 0.8B text encoder. The 6-block LLM adapter was isolated and trained to align Qwen embeddings with the DiT's latent space.
  • Stage 2: Full end-to-end fine-tuning across all 40 DiT layers with the Stage 1 adapter.

---

File Structure

├── base_model/
│   ├── UBIG-Zero.safetensors               # 40-layer DiT + 6-block llm_adapter Base Model
├── text_encoders/
│   ├── qwen35_0.8b.safetensors             # Modified Qwen 3.5 0.8B text encoder weights
└── qwen35_tokenizer/
│   ├── tokenizer.json                      # Byte-level BPE tokenizer (vocab size: 248,320)
│   └── tokenizer_config.json
└── vae/
    ├── wan21-vae.safetensors               # The original Cosmos 2 Model was trained using Wan 2.1 Vae
    └── qwen_image_vae.safetensors          # But the base model is also compatible with Qwen Image Vae

Training Details

• Framework: diffusion-pipe
• Optimizer: AdamW 8-bit (bitsandbytes)

---

Limitations & Best Practices

• Prompt Syntax: Because Qwen 3.5 is an autoregressive causal LM, it benefits from having descriptive, continuous natural language prompts.
• Precision: Run inference in bfloat16 or fp8_e4m3fn. Inference in standard fp16 can cause numeric overflow in the linear state-space recurrent scan (_delta_rule_scan).

---

Licenses

Since this isn't like some other Cosmos 2 finetuned models but rather a base raw model to build upon, I'm giving out the Base Zero Edition model into the MIT License. Whatever you build from this I guess you can try shout me out over here: https://ubtl.nx34.xyz

https://chat.nx34.xyz

---