fine_tune_lora.py

"""Task 1: Fine-tune Stable Diffusion v1.5 on Oxford Flowers-102 using LoRA.

Model selection rationale
--------------------------
Stable Diffusion v1.5 (runwayml/stable-diffusion-v1-5) was chosen because:
  - It is freely available on Hugging Face and runs on consumer GPUs.
  - The diffusers library provides a clean Python API for loading, modifying,
    and saving all sub-components (VAE, UNet, tokenizer, text encoder).
  - v1.5 is the most widely fine-tuned SD checkpoint, so LoRA tutorials and
    community knowledge transfer directly to this codebase.

Why LoRA instead of full fine-tuning
--------------------------------------
Low-Rank Adaptation (Hu et al., 2021) injects small trainable weight matrices
into the attention projections of the UNet while keeping every other parameter
frozen.  This reduces trainable parameters from ~860 M (full UNet) to ~2.3 M
(rank-4 LoRA), which:
  - Fits in 8 GB VRAM (e.g., an RTX 3050 laptop GPU).
  - Trains in minutes rather than days.
  - Avoids catastrophic forgetting of the base model's prior knowledge.

LoRA hyperparameters
---------------------
  r=4          : rank of the injected low-rank matrices; low rank = fewer
                 parameters, more regularisation.  r=4 is the standard
                 starting point for domain adaptation.
  lora_alpha=4 : scaling factor for the LoRA update (effective LR multiplier).
                 Setting alpha = r is the canonical "no scaling" choice.
  target_modules: the four attention projection layers in each transformer
                 block of the UNet (query, key, value, output projection).

Preprocessing pipeline
-----------------------
1. Resize to 512 x 512   — the VAE encoder expects this resolution.  Using a
   smaller resolution would require resampling positional embeddings.
2. CenterCrop(512)        — ensures the tensor is exactly 512x512 after resize.
3. ToTensor              — PIL uint8 -> float32 in [0, 1].
4. Normalize([0.5],[0.5]) — shifts pixel range to [-1, 1] to match the VAE's
                           expected input distribution.

Feature engineering: caption generation
----------------------------------------
Oxford Flowers-102 provides integer labels (0-101) but no text descriptions.
We map each label to a human-readable flower name and embed it in a template:

    "a high quality photo of a <flower name>, bokeh, natural light"

The quality tokens ("bokeh", "natural light") are standard positive prompts
that bias SD toward sharp, well-lit images.  Unsupported classes fall back to
"flower class N" so the code handles all 102 categories.

After training, the LoRA adapter weights are saved with unet.save_pretrained()
in Hugging Face format and can be reloaded by pipeline.py (Task 6).
"""

import os
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import Flowers102
from diffusers import StableDiffusionPipeline, DDPMScheduler
from peft import LoraConfig, get_peft_model

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
MODEL_ID   = "runwayml/stable-diffusion-v1-5"
DATA_ROOT  = "./Data"
OUT_DIR    = "./Outputs/lora"
os.makedirs(OUT_DIR, exist_ok=True)

DEVICE     = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# float16 halves VRAM usage on GPU; CPU training requires float32.
DTYPE      = torch.float16 if DEVICE.type == "cuda" else torch.float32

BATCH_SIZE = 1        # larger batches exceed 8 GB VRAM with SD; keep at 1
MAX_STEPS  = 200      # demo value — increase to 1000-2000 for a real run
LR         = 1e-4     # standard AdamW LR for LoRA fine-tuning
RESOLUTION = 512      # native SD resolution; changing this requires resampling

# ---------------------------------------------------------------------------
# Caption templates
# Human-readable names for the first 10 Flowers-102 classes.
# Classes 10-101 fall back to the generic template.
# ---------------------------------------------------------------------------
FLOWER_NAMES = [
    "pink primrose", "hard-leaved pocket orchid", "canterbury bells",
    "sweet pea", "english marigold", "tiger lily", "moon orchid",
    "bird of paradise", "monkshood", "globe thistle",
]


def caption_for(label: int) -> str:
    """Return a text prompt for a given integer class label."""
    if label < len(FLOWER_NAMES):
        return (
            f"a high quality photo of a {FLOWER_NAMES[label]}, bokeh, natural light"
        )
    return f"a high quality photo of flower class {label}, botanical, detailed petals"


# ---------------------------------------------------------------------------
# Main training loop
# ---------------------------------------------------------------------------
def main():
    print(f"Device: {DEVICE}  |  dtype: {DTYPE}")

    # Load the full SD pipeline; safety_checker is disabled to save memory.
    pipe = StableDiffusionPipeline.from_pretrained(
        MODEL_ID, torch_dtype=DTYPE, safety_checker=None
    )
    pipe.to(DEVICE)

    unet          = pipe.unet
    vae           = pipe.vae
    text_encoder  = pipe.text_encoder
    tokenizer     = pipe.tokenizer
    # DDPMScheduler adds Gaussian noise during training (forward diffusion).
    noise_scheduler = DDPMScheduler.from_config(pipe.scheduler.config)

    # Freeze VAE and text encoder — only the UNet LoRA layers will be updated.
    vae.requires_grad_(False)
    text_encoder.requires_grad_(False)
    unet.requires_grad_(False)

    # Inject LoRA adapters into the UNet attention projections.
    lora_cfg = LoraConfig(
        r=4,
        lora_alpha=4,
        target_modules=["to_q", "to_k", "to_v", "to_out.0"],
        lora_dropout=0.0,
        bias="none",
    )
    unet = get_peft_model(unet, lora_cfg)
    unet.print_trainable_parameters()  # should be ~2-3 M out of ~860 M

    optimizer = torch.optim.AdamW(
        [p for p in unet.parameters() if p.requires_grad], lr=LR
    )

    # --- Preprocessing pipeline ---------------------------------------------
    # Resize -> CenterCrop ensures a square 512x512 tensor regardless of the
    # original image aspect ratio.  Normalize maps [0,1] -> [-1,1].
    tf = transforms.Compose([
        transforms.Resize(RESOLUTION),
        transforms.CenterCrop(RESOLUTION),
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
    ])
    dataset = Flowers102(root=DATA_ROOT, split="train", download=True, transform=tf)
    loader  = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)

    # --- Training loop -------------------------------------------------------
    unet.train()
    step = 0
    for img, label in loader:
        if step >= MAX_STEPS:
            break

        img    = img.to(DEVICE, dtype=DTYPE)
        prompt = caption_for(int(label.item()))

        with torch.no_grad():
            # 1. Encode the image into the VAE latent space (4 x 64 x 64).
            latents = vae.encode(img).latent_dist.sample() * vae.config.scaling_factor

            # 2. Tokenise the caption and encode it into text embeddings.
            tok      = tokenizer(
                prompt, padding="max_length", max_length=77,
                truncation=True, return_tensors="pt",
            ).to(DEVICE)
            text_emb = text_encoder(tok.input_ids)[0]   # [1, 77, 768]

        # 3. Sample a random diffusion timestep and add the corresponding noise.
        noise      = torch.randn_like(latents)
        timesteps  = torch.randint(
            0, noise_scheduler.config.num_train_timesteps,
            (latents.size(0),), device=DEVICE,
        ).long()
        noisy      = noise_scheduler.add_noise(latents, noise, timesteps)

        # 4. Predict the noise residual (v-prediction variant uses pred_type).
        pred = unet(noisy, timesteps, encoder_hidden_states=text_emb).sample

        # 5. Denoising loss: MSE between predicted and actual noise.
        loss = F.mse_loss(pred.float(), noise.float())

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if step % 20 == 0:
            print(f"  step {step:>4}/{MAX_STEPS}  loss={loss.item():.4f}  prompt='{prompt}'")
        step += 1

    # Save LoRA adapter weights in Hugging Face PEFT format so that
    # pipeline.py can reload them with pipe.load_lora_weights().
    unet.save_pretrained(OUT_DIR)
    print(f"\nLoRA weights saved to {OUT_DIR}")


if __name__ == "__main__":
    main()