nebulette/nano-transformers

The transformer implementation of the Nano diffusion model.

The licence was kept as Apache 2.0 due to the mixed licence.

from model import Transformer2DModel
from PIL import Image
from pipeline import Pipeline
import torch
from transformers import AutoTokenizer, Gemma3ForCausalLM
# from semantic_vae import load_vae
from vae import VAE

if __name__ == '__main__':
  pipeline = Pipeline('diffusion_model.safetensors', 'unsloth/gemma-3-270m-it', 'vae.safetensors', device='cuda')
  image = pipeline.generate('1girl')

  # Manual initialization.
  vae = VAE.from_safetensors('vae.safetensors')
  # vae = load_vae('semantic_vae_step_00050000.safetensors')  # old VAE
  tokenizer = AutoTokenizer.from_pretrained('unsloth/gemma-3-270m-it')
  te = Gemma3ForCausalLM.from_pretrained('unsloth/gemma-3-270m-it')
  model = Transformer2DModel.from_safetensors('diffusion_model.safetensors')