Stable Diffusion Fine-tuning

Pre-trained text-to-image diffusion models generalize broadly but may not capture a specific visual style. Can fine-tuning only the UNet on a small, captioned Naruto-style dataset shift generation toward that aesthetic while keeping the frozen VAE and CLIP text encoder as fixed perceptual and semantic anchors?

Naruto-style dataset with text captions (Alex-0402/naruto-style-dataset-with-text via HuggingFace): 1,221 training images with paired prompts such as “a man with dark hair and brown eyes, naruto style”. Custom NarutoDataset applies resize to 512×512, RandomHorizontalFlip, ColorJitter, RandomAffine, ToTensor, and normalization to [-1, 1] with mean=std=0.5.

Base checkpoint: OFA-Sys/small-stable-diffusion-v0 (VAE + UNet2DConditionModel + CLIPTextModel + CLIPTokenizer). Only the UNet is trainable (~579M params); VAE (~84M) and text encoder (~123M) stay frozen. Training follows the standard DDPM objective: sample a timestep, add noise to VAE-encoded latents, and train the UNet to predict the noise with MSE. Inference uses StableDiffusionPipeline with PNDMScheduler, 200 steps, and classifier-free guidance scale 7.5.

AdamW on UNet only, learning rate 1e-5, batch size 1, 6 epochs on Colab GPU (7,326 steps; CUDA mixed precision with GradScaler). PNDM inference with 200 steps and classifier-free guidance 7.5. Evaluated on 5 dataset prompts, 5 custom Naruto-themed prompts, and 5 original-vs-generated pairs.

• Built NarutoDataset with augmentation beyond the assignment baseline (colour jitter, affine transforms)
• Wired VAE latent encoding/decoding with the model’s scaling_factor
• Implemented the DDPM noise-prediction training loop with CUDA mixed precision
• Froze VAE and CLIP text encoder; trained UNet weights only
• Set up StableDiffusionPipeline for text-to-image inference with classifier-free guidance
• Prepared evaluation prompts from dataset indices and custom Naruto-themed descriptions

Diffusion models feel qualitatively different from GANs—generation is iterative denoising in latent space, not a single forward pass. Fine-tuning only the UNet while freezing VAE and CLIP made the pipeline modular: perceptual compression and text semantics stay fixed while the denoiser adapts to Naruto-style captions. Completing all six epochs on Colab GPU with mixed precision produced stable training and the exported results shown here.