Upload inference.py

inference.py

@@ -0,0 +1,670 @@
+"""
+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"
+
+    @classmethod
+    def from_json(cls, path: str) -> "ShellDConfig":
+        with open(path) as f:
+            d = json.load(f)
+        # Strip keys starting with "_"
+        d = {k: v for k, v in d.items() if not k.startswith("_")}
+        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="nearest"),
+                )
+            )
+            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.Sigmoid(),
+        )
+
+    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
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn = nn.MultiheadAttention(dim, cfg.num_heads, batch_first=True)
+        self.norm_cross = nn.LayerNorm(dim)
+        self.cross_attn = nn.MultiheadAttention(dim, cfg.num_heads, batch_first=True)
+        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.Linear(dim * 4, dim),
+        )
+        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 + 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 + 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.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
+
+    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 an unconditional embedding (zeros)
+        if cfg_scale != 1.0:
+            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.")}
+
+        self.model.vae.load_state_dict(vae_sd)
+        self.model.dit.load_state_dict(dit_sd)
+
+        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
+        img_tensor = self.model.vae.decode(z)  # [1, 3, 256, 256], values in [0, 1]
+
+        # Convert to PIL
+        img_np = img_tensor[0].permute(1, 2, 0).cpu().numpy()  # [256, 256, 3]
+        img_np = (img_np * 255).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
+
+    @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]
+
+        images = []
+        for i in range(B):
+            img_np = img_tensor[i].permute(1, 2, 0).cpu().numpy()
+            img_np = (img_np * 255).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()