cds006/Progan

1# inference.py2import torch3import torch.nn.functional as F4import numpy as np5from pathlib import Path6 7from model import Generator8 9 10class ProgressiveGANInference:11    def __init__(self, checkpoint_path, device=None):12        self.device = torch.device(13            device if device is not None else14            ("cuda" if torch.cuda.is_available() else "cpu")15        )16 17        checkpoint = torch.load(checkpoint_path, map_location=self.device)18 19        # Read config safely20        if "config" in checkpoint:21            cfg = checkpoint["config"]22            self.latent_dim = cfg.get("latent_dim", 512)23            self.max_resolution = cfg.get("max_resolution", 1024)24            feature_maps = cfg.get("feature_maps", None)25        else:26            self.latent_dim = 51227            self.max_resolution = 102428            feature_maps = None29 30        self.generator = Generator(31            latent_dim=self.latent_dim,32            max_resolution=self.max_resolution,33            feature_maps=feature_maps34        ).to(self.device)35 36        # Load weights37        if "g_ema_state" in checkpoint:38            self.generator.load_state_dict(checkpoint["g_ema_state"])39        elif "generator_state" in checkpoint:40            self.generator.load_state_dict(checkpoint["generator_state"])41        else:42            raise RuntimeError("Checkpoint missing generator weights")43 44        self.generator.eval()45        self.generator.alpha = 1.046 47        self.current_resolution = checkpoint.get(48            "current_resolution", self.max_resolution49        )50        self.generator.current_resolution = self.current_resolution51 52    @torch.no_grad()53    def generate(self, num_images=1, seed=None, truncation=1.0):54        if seed is not None:55            torch.manual_seed(seed)56 57        z = torch.randn(58            num_images, self.latent_dim, 1, 1, device=self.device59        ) * truncation60 61        return self.generator(62            z,63            resolution=self.current_resolution,64            alpha=1.065        )66 67    @torch.no_grad()68    def interpolate(self, start_seed, end_seed, num_frames=8, truncation=1.0):69        torch.manual_seed(start_seed)70        z1 = torch.randn(1, self.latent_dim, 1, 1, device=self.device)71 72        torch.manual_seed(end_seed)73        z2 = torch.randn(1, self.latent_dim, 1, 1, device=self.device)74 75        z1 *= truncation76        z2 *= truncation77 78        frames = []79        for alpha in torch.linspace(0, 1, num_frames):80            z = self._slerp(z1, z2, alpha)81            img = self.generator(z, self.current_resolution, 1.0)82            frames.append(img[0])83 84        return frames85 86    @torch.no_grad()87    def generate_from_latent(self, z):88        if z.dim() == 1:89            z = z.view(1, -1, 1, 1)90        return self.generator(91            z.to(self.device),92            self.current_resolution,93            1.094        )95 96    @staticmethod97    def _slerp(z1, z2, alpha):98        z1_n = F.normalize(z1, dim=1)99        z2_n = F.normalize(z2, dim=1)100        dot = torch.clamp((z1_n * z2_n).sum(1, keepdim=True), -1, 1)101        theta = torch.acos(dot)102        sin_t = torch.sin(theta)103        return (104            torch.sin((1 - alpha) * theta) / sin_t * z1 +105            torch.sin(alpha * theta) / sin_t * z2106        )107