ShellD / inference.py
Upload inference.py
2321a0f verified 4 months ago
31.1 kB
"""
ShellD (Shell Diffusion) - Standalone Inference
================================================
Generate 256x256 images from text prompts using a pre-trained ShellD model.

This file is fully self-contained — it does NOT import from train.py.
It duplicates only the architecture classes needed for inference, with all
training-specific code (dataset, optimizers, VAE pretraining, git cloning, etc.) removed.

Usage (load from Hugging Face):
    from inference import ShellDInference

    pipe = ShellDInference("FlameF0X/ShellD")
    img = pipe.generate("a mountain lake at sunset")
    img.save("output.png")

Usage (load from local directory):
    pipe = ShellDInference("./ShellD_model")
"""

import json
import math
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Dict, Any, List, Optional, Tuple

import torch
import torch.nn as nn
from huggingface_hub import snapshot_download
import torch.nn.functional as F
from safetensors.torch import load_file
from sentence_transformers import SentenceTransformer
from PIL import Image
import numpy as np


# ═══════════════════════════════════════════════════════════
# Config
# ═══════════════════════════════════════════════════════════

@dataclass
class ShellDConfig:
    """Model hyperparameters. Must match the saved config.json."""
    image_size: int = 256
    latent_dim: int = 16
    ae_hidden_dim: int = 64
    ae_num_blocks: int = 3
    hidden_dim: int = 256
    num_hidden_layers: int = 12
    num_heads: int = 8
    patch_size: int = 4
    text_encoder_name: str = "./all-MiniLM-L6-v2"
    text_encoder_dim: int = 384
    num_timesteps: int = 1000
    beta_start: float = 1e-4
    beta_end: float = 0.02
    model_name: str = "ShellD"
    dropout: float = 0.0               # match train.py; 0 = no dropout (inference)
    kl_weight: float = 0.1             # VAE KL weight (not used in inference)
    cfg_dropout_prob: float = 0.15     # CFG dropout prob (not used in inference)

    @classmethod
    def from_json(cls, path: str) -> "ShellDConfig":
        with open(path) as f:
            d = json.load(f)
        # Keep only fields that exist in the dataclass
        valid_keys = {f.name for f in cls.__dataclass_fields__.values()}
        d = {k: v for k, v in d.items() if k in valid_keys}
        return cls(**d)


# ═══════════════════════════════════════════════════════════
# VAE — Encoder / Decoder
# ═══════════════════════════════════════════════════════════

class ResidualBlock(nn.Module):
    def __init__(self, in_ch: int, out_ch: int):
        super().__init__()
        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.norm1 = nn.GroupNorm(8, out_ch)
        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
        self.norm2 = nn.GroupNorm(8, out_ch)
        self.skip = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = self.skip(x)
        x = F.silu(self.norm1(self.conv1(x)))
        x = self.norm2(self.conv2(x))
        return F.silu(x + residual)


class Encoder(nn.Module):
    def __init__(self, cfg: ShellDConfig):
        super().__init__()
        self.cfg = cfg
        self.init_conv = nn.Conv2d(3, cfg.ae_hidden_dim, 3, padding=1)

        blocks = []
        ch = cfg.ae_hidden_dim
        for _ in range(cfg.ae_num_blocks):
            out_ch = min(ch * 2, 256)
            blocks.append(
                nn.Sequential(
                    ResidualBlock(ch, out_ch),
                    ResidualBlock(out_ch, out_ch),
                    nn.Conv2d(out_ch, out_ch, 4, stride=2, padding=1),  # downsample
                )
            )
            ch = out_ch
        self.down_blocks = nn.ModuleList(blocks)

        self.mid = nn.Sequential(
            ResidualBlock(ch, ch),
            ResidualBlock(ch, ch),
        )
        self.out_conv = nn.Conv2d(ch, cfg.latent_dim * 2, 3, padding=1)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        x = self.init_conv(x)
        for block in self.down_blocks:
            x = block(x)
        x = self.mid(x)
        x = self.out_conv(x)
        mean, logvar = x.chunk(2, dim=1)
        return mean, logvar


class Decoder(nn.Module):
    def __init__(self, cfg: ShellDConfig):
        super().__init__()
        self.cfg = cfg
        ch = cfg.ae_hidden_dim * (2 ** cfg.ae_num_blocks)
        self.init_conv = nn.Conv2d(cfg.latent_dim, ch, 3, padding=1)

        self.mid = nn.Sequential(
            ResidualBlock(ch, ch),
            ResidualBlock(ch, ch),
        )

        up_blocks = []
        for _ in range(cfg.ae_num_blocks):
            out_ch = ch // 2
            up_blocks.append(
                nn.Sequential(
                    ResidualBlock(ch, out_ch),
                    ResidualBlock(out_ch, out_ch),
                    nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
                )
            )
            ch = out_ch
        self.up_blocks = nn.ModuleList(up_blocks)

        self.out_conv = nn.Sequential(
            ResidualBlock(ch, ch),
            nn.Conv2d(ch, 3, 3, padding=1),
            nn.Tanh(),
        )

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        x = self.init_conv(z)
        x = self.mid(x)
        for block in self.up_blocks:
            x = block(x)
        return self.out_conv(x)


class VAE(nn.Module):
    def __init__(self, cfg: ShellDConfig):
        super().__init__()
        self.cfg = cfg
        self.encoder = Encoder(cfg)
        self.decoder = Decoder(cfg)

    def encode(self, x: torch.Tensor) -> torch.Tensor:
        mean, logvar = self.encoder(x)
        logvar = logvar.clamp(-10.0, 10.0)
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mean + eps * std

    def decode(self, z: torch.Tensor) -> torch.Tensor:
        return self.decoder(z)

    def forward(self, x: torch.Tensor):
        """Full forward (used only for training; included for completeness)."""
        mean, logvar = self.encoder(x)
        logvar = logvar.clamp(-10.0, 10.0)
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        z = mean + eps * std
        return self.decode(z)


# ═══════════════════════════════════════════════════════════
# DiT Backbone
# ═══════════════════════════════════════════════════════════

class PatchEmbed(nn.Module):
    def __init__(self, cfg: ShellDConfig):
        super().__init__()
        self.cfg = cfg
        self.latent_size = cfg.image_size // (2 ** cfg.ae_num_blocks)
        self.num_patches_1d = self.latent_size // cfg.patch_size
        self.num_patches = self.num_patches_1d ** 2
        self.patch_dim = cfg.latent_dim * (cfg.patch_size ** 2)
        self.proj = nn.Linear(self.patch_dim, cfg.hidden_dim)

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        B, C, H, W = z.shape
        ps = self.cfg.patch_size
        n = self.num_patches_1d
        z = z.reshape(B, C, n, ps, n, ps)
        z = z.permute(0, 2, 4, 1, 3, 5).reshape(B, self.num_patches, self.patch_dim)
        return self.proj(z)


class DiTBlock(nn.Module):
    def __init__(self, cfg: ShellDConfig):
        super().__init__()
        dim = cfg.hidden_dim
        drop = cfg.dropout

        self.norm1 = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, cfg.num_heads, batch_first=True)
        self.drop_attn = nn.Dropout(drop)

        self.norm_cross = nn.LayerNorm(dim)
        self.cross_attn = nn.MultiheadAttention(dim, cfg.num_heads, batch_first=True)
        self.drop_cross = nn.Dropout(drop)
        self.text_proj = nn.Linear(cfg.text_encoder_dim, dim)

        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Dropout(drop),
            nn.Linear(dim * 4, dim),
        )
        self.drop_mlp = nn.Dropout(drop)

        self.timestep_mlp = nn.Sequential(
            nn.Linear(dim, dim * 4),
            nn.SiLU(),
            nn.Linear(dim * 4, dim),
        )

    def forward(self, x: torch.Tensor, text_emb: torch.Tensor, t_emb: torch.Tensor):
        # Self-attention
        h = self.norm1(x)
        attn_out, _ = self.attn(h, h, h)
        x = x + self.drop_attn(attn_out)
        # Cross-attention with text
        text_proj = self.text_proj(text_emb)
        h = self.norm_cross(x)
        cross_out, _ = self.cross_attn(h, text_proj, text_proj)
        x = x + self.drop_cross(cross_out)
        # Timestep conditioning
        t_proj = self.timestep_mlp(t_emb)
        x = x + t_proj.unsqueeze(1)
        # MLP
        h = self.norm2(x)
        x = x + self.drop_mlp(self.mlp(h))
        return x


class DiT(nn.Module):
    def __init__(self, cfg: ShellDConfig):
        super().__init__()
        self.cfg = cfg
        self.latent_size = cfg.image_size // (2 ** cfg.ae_num_blocks)
        self.num_patches_1d = self.latent_size // cfg.patch_size
        self.num_patches = self.num_patches_1d ** 2

        self.patch_embed = PatchEmbed(cfg)
        self.pos_embed = nn.Parameter(torch.randn(1, self.num_patches, cfg.hidden_dim))
        self.blocks = nn.ModuleList([DiTBlock(cfg) for _ in range(cfg.num_hidden_layers)])
        self.norm = nn.LayerNorm(cfg.hidden_dim)
        self.out_proj = nn.Linear(cfg.hidden_dim, cfg.latent_dim * (cfg.patch_size ** 2))

        self.time_mlp = nn.Sequential(
            nn.Linear(cfg.hidden_dim, cfg.hidden_dim * 4),
            nn.SiLU(),
            nn.Linear(cfg.hidden_dim * 4, cfg.hidden_dim),
        )

    @staticmethod
    def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000):
        half = dim // 2
        freqs = torch.exp(
            -math.log(max_period) * torch.arange(0, half, dtype=torch.float32) / half
        ).to(t.device)
        args = t[:, None].float() * freqs[None]
        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        if dim % 2 == 1:
            emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
        return emb

    def forward(self, z: torch.Tensor, text_emb: torch.Tensor, t: torch.Tensor):
        B = z.shape[0]
        dim = self.cfg.hidden_dim
        t_emb = self.timestep_embedding(t, dim)
        t_emb = self.time_mlp(t_emb)

        x = self.patch_embed(z)
        x = x + self.pos_embed

        for blk in self.blocks:
            x = blk(x, text_emb, t_emb)

        x = self.norm(x)
        x = self.out_proj(x)

        ps = self.cfg.patch_size
        H = W = self.latent_size
        n = self.num_patches_1d
        x = x.reshape(B, n, n, self.cfg.latent_dim, ps, ps)
        x = x.permute(0, 3, 1, 4, 2, 5).reshape(B, self.cfg.latent_dim, H, W)
        return x


# ═══════════════════════════════════════════════════════════
# Full ShellD Model (Inference-only)
# ═══════════════════════════════════════════════════════════

class ShellDModel(nn.Module):
    """ShellD model for inference. Minimal — no training helpers."""

    def __init__(self, cfg: ShellDConfig):
        super().__init__()
        self.cfg = cfg
        self.vae = VAE(cfg)
        self.dit = DiT(cfg)
        self.text_encoder = None  # loaded separately
        # Null text embedding for classifier-free guidance (loaded from checkpoint)
        self.null_text_embed: Optional[torch.Tensor] = None

    def encode_text(self, prompts: List[str], device: torch.device) -> torch.Tensor:
        assert self.text_encoder is not None, "Text encoder not loaded"
        with torch.no_grad():
            emb = self.text_encoder.encode(prompts, convert_to_tensor=True)
        return emb.to(device).unsqueeze(1)  # [B, 1, 384]

    def forward(self, z: torch.Tensor, text_emb: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """Predict the noise at timestep t given noisy latent z and text."""
        return self.dit(z, text_emb, t)


# ═══════════════════════════════════════════════════════════
# Diffusion Schedule (Reverse Process Only)
# ═══════════════════════════════════════════════════════════

class DiffusionSchedule:
    """DDPM schedule for the reverse (denoising) process."""

    def __init__(self, cfg: ShellDConfig, device: torch.device):
        betas = torch.linspace(cfg.beta_start, cfg.beta_end, cfg.num_timesteps, device=device)
        alphas = 1.0 - betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)

        self.betas = betas
        self.alphas = alphas
        self.alphas_cumprod = alphas_cumprod
        self.sqrt_alphas_cumprod = alphas_cumprod.sqrt()
        self.sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod).sqrt()

    @torch.no_grad()
    def sample(
        self,
        model: ShellDModel,
        text_emb: torch.Tensor,
        num_steps: Optional[int] = None,
        cfg_scale: float = 3.0,
        seed: Optional[int] = None,
    ) -> torch.Tensor:
        """
        DDPM reverse sampling with optional classifier-free guidance.

        Args:
            model: The ShellD model.
            text_emb: Text embedding [B, 1, 384].
            num_steps: Number of denoising steps (default: cfg.num_timesteps).
            cfg_scale: Classifier-free guidance scale. 1.0 = no guidance.
            seed: Optional random seed for reproducibility.

        Returns:
            Denoised latent tensor [B, latent_dim, H', W'].
        """
        if seed is not None:
            torch.manual_seed(seed)

        device = text_emb.device
        B = text_emb.shape[0]
        cfg = model.cfg

        num_steps = num_steps or cfg.num_timesteps

        # Latent spatial size
        H = W = cfg.image_size // (2 ** cfg.ae_num_blocks)

        # Start from random noise
        z = torch.randn(B, cfg.latent_dim, H, W, device=device)

        # For classifier-free guidance, we need the learned null embedding
        if cfg_scale != 1.0:
            if model.null_text_embed is not None:
                uncond_emb = model.null_text_embed.to(device).expand(B, -1, -1)
            else:
                # Fallback: zeros (works if model was trained with zero-dropped text)
                uncond_emb = torch.zeros_like(text_emb)

        # Time step resampling for faster inference (DDPM-style, evenly spaced)
        step_indices = torch.linspace(0, cfg.num_timesteps - 1, num_steps, device=device, dtype=torch.long)

        for i in range(num_steps - 1, -1, -1):
            t = step_indices[i]
            t_batch = t.expand(B)

            # Predict noise
            if cfg_scale != 1.0:
                # Classifier-free guidance: combine conditional and unconditional predictions
                z_in = torch.cat([z, z], dim=0)
                t_in = torch.cat([t_batch, t_batch], dim=0)
                text_in = torch.cat([text_emb, uncond_emb], dim=0)
                noise_pred = model(z_in, text_in, t_in)
                noise_cond, noise_uncond = noise_pred.chunk(2, dim=0)
                noise_pred = noise_uncond + cfg_scale * (noise_cond - noise_uncond)
            else:
                noise_pred = model(z, text_emb, t_batch)

            # DDPM update step
            alpha = self.alphas[t]
            alpha_cumprod = self.alphas_cumprod[t]
            beta = self.betas[t]

            # Compute predicted x0 (for logging purposes, not needed for step)
            # x0_pred = (z - sqrt_one_minus_ac * noise_pred) / sqrt_ac

            # Mean for posterior
            coef1 = 1.0 / alpha.sqrt()
            coef2 = beta / (1.0 - alpha_cumprod).sqrt()
            z_mean = coef1 * (z - coef2 * noise_pred)

            if i > 0:
                noise = torch.randn_like(z)
                z = z_mean + beta.sqrt() * noise
            else:
                z = z_mean

        return z


# ═══════════════════════════════════════════════════════════
# High-Level Inference Pipeline
# ═══════════════════════════════════════════════════════════

class ShellDInference:
    """
    High-level pipeline for ShellD text-to-image generation.

    Usage:
        pipe = ShellDInference("./ShellD_model")
        img = pipe.generate("a cat sitting on a mat")
        img.save("cat.png")
    """

    def __init__(
        self,
        model_dir: str,
        device: Optional[str] = None,
        text_encoder_path: Optional[str] = None,
    ):
        """
        Args:
            model_dir: Path to the model directory, or a Hugging Face repo ID
                       (e.g. "FlameF0X/ShellD"). If a repo ID is given, the
                       weights are automatically downloaded via huggingface_hub.
            device: Device to run on ('cuda', 'cpu', or None for auto-detect).
            text_encoder_path: Optional override path for the text encoder model.
                               Defaults to the path in config.json.
        """
        self.device = torch.device(
            device or ("cuda" if torch.cuda.is_available() else "cpu")
        )
        print(f"ShellD inference using device: {self.device}")

        # Resolve model source: Hugging Face repo or local path
        local_path = Path(model_dir)
        if local_path.exists():
            # Local directory
            self.model_dir = local_path
        else:
            # Assume it's a Hugging Face repo ID — download via hub
            print(f"Downloading model from Hugging Face: {model_dir}")
            self.model_dir = Path(snapshot_download(repo_id=model_dir))
            print(f"Model cached at: {self.model_dir}")

        # 1. Load config
        config_path = self.model_dir / "config.json"
        if not config_path.exists():
            raise FileNotFoundError(f"config.json not found at {config_path}")
        self.cfg = ShellDConfig.from_json(str(config_path))

        # 2. Build model
        self.model = ShellDModel(self.cfg).to(self.device)
        self.model.eval()

        # 3. Load weights
        weights_path = self.model_dir / "model.safetensors"
        if not weights_path.exists():
            raise FileNotFoundError(f"model.safetensors not found at {weights_path}")
        self._load_weights(str(weights_path))

        # 4. Load text encoder
        encoder_path = text_encoder_path or self.cfg.text_encoder_name
        print(f"Loading text encoder from: {encoder_path}")
        self.model.text_encoder = SentenceTransformer(encoder_path)

        # 5. Diffusion schedule
        self.diffusion = DiffusionSchedule(self.cfg, self.device)

        # Print parameter count
        total = sum(p.numel() for p in self.model.parameters())
        trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        print(f"Model loaded: {total:,} total params ({trainable:,} trainable)")

    def _load_weights(self, weights_path: str):
        """Load state dict, mapping prefixes to the correct submodules."""
        sd = load_file(weights_path)

        vae_sd = {k.replace("vae.", ""): v for k, v in sd.items() if k.startswith("vae.")}
        dit_sd = {k.replace("dit.", ""): v for k, v in sd.items() if k.startswith("dit.")}
        txt_sd = {k.replace("text_encoder.", ""): v for k, v in sd.items() if k.startswith("text_encoder.")}

        # Backward compat: old checkpoints (no dropout) have final Linear at mlp.2;
        # new architecture (with dropout) has it at mlp.3. Remap silently.
        if any(".mlp.2.weight" in k for k in dit_sd) and not any(".mlp.3.weight" in k for k in dit_sd):
            remapped = {}
            for k, v in dit_sd.items():
                remapped[k.replace(".mlp.2.", ".mlp.3.")] = v
            dit_sd = remapped
            print("  ↻ Remapped old DiT checkpoint (no dropout) → new MLP layout.")

        self.model.vae.load_state_dict(vae_sd)
        self.model.dit.load_state_dict(dit_sd)

        # Restore null text embedding for CFG
        if "null_text_embed" in sd:
            self.model.null_text_embed = sd["null_text_embed"].to(self.device)
            print("Null text embedding loaded for CFG.")
        else:
            print("No null_text_embed in checkpoint — CFG will use zeros (may degrade quality).")

        if txt_sd and self.model.text_encoder is not None:
            self.model.text_encoder.load_state_dict(txt_sd)
            print("Text encoder weights loaded from safetensors.")
        else:
            print("Text encoder loaded from SentenceTransformer cache (weights not in safetensors).")

    @torch.no_grad()
    def generate(
        self,
        prompt: str,
        num_steps: int = 250,
        cfg_scale: float = 3.0,
        seed: Optional[int] = None,
        output_size: Optional[int] = None,
    ) -> Image.Image:
        """
        Generate an image from a text prompt.

        Args:
            prompt: Text description of the desired image.
            num_steps: Number of denoising steps (fewer = faster, lower quality).
                       Recommended: 250-1000.
            cfg_scale: Classifier-free guidance scale. Higher = more prompt adherence.
                       1.0 = no guidance. Typical range: 2.0-5.0.
            seed: Random seed for reproducibility.
            output_size: If set, the output image is resized to (output_size, output_size).

        Returns:
            A PIL Image.
        """
        self.model.eval()

        # Encode prompt
        text_emb = self.model.encode_text([prompt], self.device)  # [1, 1, 384]

        # Sample latent
        z = self.diffusion.sample(
            self.model,
            text_emb,
            num_steps=num_steps,
            cfg_scale=cfg_scale,
            seed=seed,
        )

        # Decode latent to image (VAE decoder outputs [-1, 1] via Tanh)
        img_tensor = self.model.vae.decode(z)  # [1, 3, 256, 256], values in [-1, 1]

        # Convert to PIL: [-1,1] → [0,255]
        img_np = img_tensor[0].permute(1, 2, 0).cpu().numpy()  # [256, 256, 3]
        img_np = ((img_np + 1.0) * 127.5).clip(0, 255).astype(np.uint8)
        img = Image.fromarray(img_np)

        if output_size is not None:
            img = img.resize((output_size, output_size), Image.Resampling.BICUBIC)

        return img

    def _decode_latent(self, z: torch.Tensor) -> Image.Image:
        """Decode a latent tensor to a PIL Image (helper for streaming)."""
        img_tensor = self.model.vae.decode(z)  # [B, 3, H, W] in [-1,1]
        img_np = img_tensor[0].permute(1, 2, 0).cpu().numpy()
        img_np = ((img_np + 1.0) * 127.5).clip(0, 255).astype(np.uint8)
        return Image.fromarray(img_np)

    @torch.no_grad()
    def generate_stream(
        self,
        prompt: str,
        num_steps: int = 250,
        cfg_scale: float = 3.0,
        seed: Optional[int] = None,
        display_every: int = 10,
    ):
        """
        Generator that yields (PIL.Image, str) tuples showing the denoising
        process unfold step-by-step.  Callers can display each intermediate
        image as it's produced for a live "diffusion reveal" effect.

        Yields:
            (image, status_label) — status_label is e.g. "Step 50/250".
        """
        self.model.eval()
        device = self.device
        cfg = self.model.cfg

        if seed is not None:
            torch.manual_seed(seed)

        # --- Encode prompt ---
        text_emb = self.model.encode_text([prompt], device)  # [1, 1, 384]
        if cfg_scale != 1.0:
            if self.model.null_text_embed is not None:
                uncond_emb = self.model.null_text_embed.to(device).expand(1, -1, -1)
            else:
                uncond_emb = torch.zeros_like(text_emb)
        else:
            uncond_emb = None

        # --- Start from pure noise ---
        H = W = cfg.image_size // (2 ** cfg.ae_num_blocks)
        z = torch.randn(1, cfg.latent_dim, H, W, device=device)

        # --- Timestep schedule (evenly spaced) ---
        step_indices = torch.linspace(
            0, cfg.num_timesteps - 1, num_steps, device=device, dtype=torch.long
        )

        # Yield initial noise snapshot
        yield self._decode_latent(z), f"Step 0/{num_steps} — pure noise"

        # --- DDPM reverse denoising loop ---
        for i in range(num_steps - 1, -1, -1):
            t = step_indices[i]
            t_batch = t.expand(1)

            # Predict noise (with classifier-free guidance)
            if cfg_scale != 1.0:
                z_in = torch.cat([z, z], dim=0)
                t_in = torch.cat([t_batch, t_batch], dim=0)
                text_in = torch.cat([text_emb, uncond_emb], dim=0)
                noise_pred = self.model(z_in, text_in, t_in)
                noise_cond, noise_uncond = noise_pred.chunk(2, dim=0)
                noise_pred = noise_uncond + cfg_scale * (noise_cond - noise_uncond)
            else:
                noise_pred = self.model(z, text_emb, t_batch)

            # DDPM update
            alpha = self.diffusion.alphas[t]
            alpha_cumprod = self.diffusion.alphas_cumprod[t]
            beta = self.diffusion.betas[t]

            coef1 = 1.0 / alpha.sqrt()
            coef2 = beta / (1.0 - alpha_cumprod).sqrt()
            z_mean = coef1 * (z - coef2 * noise_pred)

            if i > 0:
                z = z_mean + beta.sqrt() * torch.randn_like(z)
            else:
                z = z_mean

            # Yield intermediate snapshot at display intervals
            step_no = num_steps - i
            if step_no % display_every == 0 or i == 0:
                yield self._decode_latent(z), f"Step {step_no}/{num_steps}"

    @torch.no_grad()
    def generate_batch(
        self,
        prompts: List[str],
        num_steps: int = 250,
        cfg_scale: float = 3.0,
        seed: Optional[int] = None,
    ) -> List[Image.Image]:
        """
        Generate images for multiple prompts efficiently (batched).

        Args:
            prompts: List of text prompts.
            num_steps: Number of denoising steps.
            cfg_scale: Classifier-free guidance scale.
            seed: Random seed (applied per-prompt).

        Returns:
            List of PIL Images.
        """
        self.model.eval()
        B = len(prompts)

        # Encode all prompts
        text_emb = self.model.encode_text(prompts, self.device)  # [B, 1, 384]

        # Sample latents
        z = self.diffusion.sample(
            self.model,
            text_emb,
            num_steps=num_steps,
            cfg_scale=cfg_scale,
            seed=seed,
        )

        # Decode
        img_tensor = self.model.vae.decode(z)  # [B, 3, 256, 256] in [-1,1]

        images = []
        for i in range(B):
            img_np = img_tensor[i].permute(1, 2, 0).cpu().numpy()
            img_np = ((img_np + 1.0) * 127.5).clip(0, 255).astype(np.uint8)
            images.append(Image.fromarray(img_np))

        return images


# ═══════════════════════════════════════════════════════════
# Command-Line Interface
# ═══════════════════════════════════════════════════════════

def main():
    import argparse

    parser = argparse.ArgumentParser(description="ShellD - Text-to-Image Generation")
    parser.add_argument("--model_dir", type=str, default="FlameF0X/ShellD",
                        help="Path to model directory or Hugging Face repo ID (default: FlameF0X/ShellD)")
    parser.add_argument("--text_encoder", type=str, default=None,
                        help="Override path for the text encoder model")
    parser.add_argument("--prompt", type=str, required=True,
                        help="Text prompt for image generation")
    parser.add_argument("--output", type=str, default="output.png",
                        help="Output image path")
    parser.add_argument("--steps", type=int, default=250,
                        help="Number of denoising steps (default: 250)")
    parser.add_argument("--cfg", type=float, default=3.0,
                        help="Classifier-free guidance scale (default: 3.0)")
    parser.add_argument("--seed", type=int, default=None,
                        help="Random seed for reproducibility")
    parser.add_argument("--device", type=str, default=None,
                        help="Device: 'cuda' or 'cpu'")
    parser.add_argument("--size", type=int, default=None,
                        help="Output image size (square, resized)")

    args = parser.parse_args()

    pipe = ShellDInference(
        model_dir=args.model_dir,
        device=args.device,
        text_encoder_path=args.text_encoder,
    )

    img = pipe.generate(
        prompt=args.prompt,
        num_steps=args.steps,
        cfg_scale=args.cfg,
        seed=args.seed,
        output_size=args.size,
    )

    img.save(args.output)
    print(f"Image saved to {args.output}")


if __name__ == "__main__":
    main()