procgne/Plonk

1import torch2import random3import string4from transformers import AutoTokenizer, T5EncoderModel5from models.pretrained_models import Plonk6from models.samplers.riemannian_flow_sampler import riemannian_flow_sampler7 8from models.postprocessing import CartesiantoGPS9 10from models.schedulers import (11    SigmoidScheduler,12    LinearScheduler,13    CosineScheduler,14)15from models.preconditioning import DDPMPrecond16from torchvision import transforms17from transformers import CLIPProcessor, CLIPVisionModel18from utils.image_processing import CenterCrop19import numpy as np20 21device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")22MODELS = {23    "nicolas-dufour/PLONK_YFCC": {"emb_name": "dinov2"},24    "nicolas-dufour/PLONK_OSV_5M": {25        "emb_name": "street_clip",26    },27    "nicolas-dufour/PLONK_iNaturalist": {28        "emb_name": "dinov2",29    },30}31 32 33def scheduler_fn(34    scheduler_type: str, start: float, end: float, tau: float, clip_min: float = 1e-935):36    if scheduler_type == "sigmoid":37        return SigmoidScheduler(start, end, tau, clip_min)38    elif scheduler_type == "cosine":39        return CosineScheduler(start, end, tau, clip_min)40    elif scheduler_type == "linear":41        return LinearScheduler(clip_min=clip_min)42    else:43        raise ValueError(f"Scheduler type {scheduler_type} not supported")44 45 46class DinoV2FeatureExtractor:47    def __init__(self, device=device):48        super().__init__()49        self.device = device50        self.emb_model = torch.hub.load("facebookresearch/dinov2", "dinov2_vitl14_reg")51        self.emb_model.eval()52        self.emb_model.to(self.device)53        self.augmentation = transforms.Compose(54            [55                CenterCrop(ratio="1:1"),56                transforms.Resize(57                    336, interpolation=transforms.InterpolationMode.BICUBIC58                ),59                transforms.ToTensor(),60                transforms.Normalize(61                    mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)62                ),63            ]64        )65 66    def __call__(self, batch):67        embs = []68        with torch.no_grad():69            for img in batch["img"]:70                emb = self.emb_model(71                    self.augmentation(img).unsqueeze(0).to(self.device)72                ).squeeze(0)73                embs.append(emb)74        batch["emb"] = torch.stack(embs)75        return batch76 77 78class StreetClipFeatureExtractor:79    def __init__(self, device=device):80        self.device = device81        self.emb_model = CLIPVisionModel.from_pretrained("geolocal/StreetCLIP").to(82            device83        )84        self.processor = CLIPProcessor.from_pretrained("geolocal/StreetCLIP")85 86    def __call__(self, batch):87        inputs = self.processor(images=batch["img"], return_tensors="pt")88        inputs = {k: v.to(self.device) for k, v in inputs.items()}89        with torch.no_grad():90            outputs = self.emb_model(**inputs)91        embeddings = outputs.last_hidden_state[:, 0]92        batch["emb"] = embeddings93        return batch94 95 96def load_prepocessing(model_name, dtype=torch.float32):97    if MODELS[model_name]["emb_name"] == "dinov2":98        return DinoV2FeatureExtractor()99    elif MODELS[model_name]["emb_name"] == "street_clip":100        return StreetClipFeatureExtractor()101    else:102        raise ValueError(f"Embedding model {MODELS[model_name]['emb_name']} not found")103 104 105class PlonkPipeline:106    """107    The CADT2IPipeline class is designed to facilitate the generation of images from text prompts using a pre-trained CAD model.108    It integrates various components such as samplers, schedulers, and post-processing techniques to produce high-quality images.109 110    Initialization:111        CADT2IPipeline(112            model_path,113            sampler="ddim",114            scheduler="sigmoid",115            postprocessing="sd_1_5_vae",116            scheduler_start=-3,117            scheduler_end=3,118            scheduler_tau=1.1,119            device="cuda",120        )121 122    Parameters:123        model_path (str): Path to the pre-trained CAD model.124        sampler (str): The sampling method to use. Options are "ddim", "ddpm", "dpm", "dpm_2S", "dpm_2M". Default is "ddim".125        scheduler (str): The scheduler type to use. Options are "sigmoid", "cosine", "linear". Default is "sigmoid".126        postprocessing (str): The post-processing method to use. Options are "consistency-decoder", "sd_1_5_vae". Default is "sd_1_5_vae".127        scheduler_start (float): Start value for the scheduler. Default is -3.128        scheduler_end (float): End value for the scheduler. Default is 3.129        scheduler_tau (float): Tau value for the scheduler. Default is 1.1.130        device (str): Device to run the model on. Default is "cuda".131 132    Methods:133        model(*args, **kwargs):134            Runs the preconditioning on the network with the provided arguments.135 136        __call__(...):137            Generates images based on the provided conditions and parameters.138 139            Parameters:140                cond (str or list of str): The conditioning text or list of texts.141                num_samples (int, optional): Number of samples to generate. If not provided, it is inferred from cond.142                x_N (torch.Tensor, optional): Initial noise tensor. If not provided, it is generated.143                latents (torch.Tensor, optional): Previous latents.144                num_steps (int, optional): Number of steps for the sampler. If not provided, the default is used.145                sampler (callable, optional): Custom sampler function. If not provided, the default sampler is used.146                scheduler (callable, optional): Custom scheduler function. If not provided, the default scheduler is used.147                cfg (float): Classifier-free guidance scale. Default is 15.148                guidance_type (str): Type of guidance. Default is "constant".149                guidance_start_step (int): Step to start guidance. Default is 0.150                generator (torch.Generator, optional): Random number generator.151                coherence_value (float): Doherence value for sampling. Default is 1.0.152                uncoherence_value (float): Uncoherence value for sampling. Default is 0.0.153                unconfident_prompt (str, optional): Unconfident prompt text.154                thresholding_type (str): Type of thresholding. Default is "clamp".155                clamp_value (float): Clamp value for thresholding. Default is 1.0.156                thresholding_percentile (float): Percentile for thresholding. Default is 0.995.157 158            Returns:159                torch.Tensor: The generated image tensor after post-processing.160 161        to(device):162            Moves the model and its components to the specified device.163 164            Parameters:165                device (str): The device to move the model to (e.g., "cuda", "cpu").166 167            Returns:168                CADT2IPipeline: The pipeline instance with updated device.169 170    Example Usage:171            pipe = CADT2IPipeline(172                "nicolas-dufour/",173            )174            pipe.to("cuda")175            image = pipe(176                "a beautiful landscape with a river and mountains",177                num_samples=4,178            )179    """180 181    def __init__(182        self,183        model_path,184        scheduler="sigmoid",185        scheduler_start=-7,186        scheduler_end=3,187        scheduler_tau=1.0,188        device=device,189    ):190        self.network = Plonk.from_pretrained(model_path).to(device)191        self.network.requires_grad_(False).eval()192        assert scheduler in [193            "sigmoid",194            "cosine",195            "linear",196        ], f"Scheduler {scheduler} not supported"197        self.scheduler = scheduler_fn(198            scheduler, scheduler_start, scheduler_end, scheduler_tau199        )200        self.cond_preprocessing = load_prepocessing(model_name=model_path)201        self.postprocessing = CartesiantoGPS()202        self.sampler = riemannian_flow_sampler203        self.model_path = model_path204        self.preconditioning = DDPMPrecond()205        self.device = device206 207    def model(self, *args, **kwargs):208        return self.preconditioning(self.network, *args, **kwargs)209 210    def __call__(211        self,212        images,213        batch_size=None,214        x_N=None,215        num_steps=None,216        scheduler=None,217        cfg=0,218        generator=None,219        callback=None,220    ):221        """Sample from the model given conditioning.222 223        Args:224            cond: Conditioning input (image or list of images)225            batch_size: Number of samples to generate (inferred from cond if not provided)226            x_N: Initial noise tensor (generated if not provided)227            num_steps: Number of sampling steps (uses default if not provided)228            sampler: Custom sampler function (uses default if not provided)229            scheduler: Custom scheduler function (uses default if not provided)230            cfg: Classifier-free guidance scale (default 15)231            generator: Random number generator232            callback: Optional callback function to report progress (step, total_steps)233 234        Returns:235            Sampled GPS coordinates after postprocessing236        """237        # Set up batch size and initial noise238        shape = [3]239        if not isinstance(images, list):240            images = [images]241        if x_N is None:242            if batch_size is None:243                if isinstance(images, list):244                    batch_size = len(images)245                else:246                    batch_size = 1247            x_N = torch.randn(248                batch_size, *shape, device=self.device, generator=generator249            )250        else:251            x_N = x_N.to(self.device)252            if x_N.ndim == 3:253                x_N = x_N.unsqueeze(0)254            batch_size = x_N.shape[0]255 256        # Set up batch with conditioning257        batch = {"y": x_N}258        batch["img"] = images259        batch = self.cond_preprocessing(batch)260        if len(images) > 1:261            assert len(images) == batch_size262        else:263            batch["emb"] = batch["emb"].repeat(batch_size, 1)264 265        # Use default sampler/scheduler if not provided266        sampler = self.sampler267        if scheduler is None:268            scheduler = self.scheduler269 270        # Sample from model271        if num_steps is None:272            num_steps = 16  # Default number of steps273 274        # Create a wrapper for the model that updates progress275        def model_with_progress(*args, **kwargs):276            step = kwargs.pop('current_step', 0)277            if callback:278                callback(step, num_steps)279            return self.model(*args, **kwargs)280 281        output = sampler(282            model_with_progress,283            batch,284            conditioning_keys="emb",285            scheduler=scheduler,286            num_steps=num_steps,287            cfg_rate=cfg,288            generator=generator,289            callback=callback,290        )291 292        # Apply postprocessing and return293        output = self.postprocessing(output)294        # To degrees295        output = np.degrees(output.detach().cpu().numpy())296        return output297 298    def to(self, device):299        self.network.to(device)300        self.postprocessing.to(device)301        self.device = torch.device(device)302        return self303