Buckets:
| """GPU model bundle for offline rollout decoding and visual quality metrics.""" | |
| from __future__ import annotations | |
| import math | |
| import torch | |
| import torch.nn.functional as F | |
| from PIL import Image | |
| from .visual_quality import align_generated_to_reference, finite_number | |
| _MS_SSIM_WEIGHTS = (0.0448, 0.2856, 0.3001, 0.2363, 0.1333) | |
| class RolloutRGBDecoder: | |
| """Batch same-grid BSQ states through the released UniAR RGB decoder.""" | |
| def __init__(self, model_path, device: torch.device): | |
| from uniar import UniARVisualDecoder | |
| self.device = device | |
| self.decoder = UniARVisualDecoder.from_pretrained(model_path, device=device) | |
| def decode( | |
| self, | |
| codes: list[torch.Tensor], | |
| *, | |
| grid: tuple[int, int], | |
| generators: list[torch.Generator], | |
| steps: int, | |
| guidance: float, | |
| upsampling: float, | |
| ) -> list[Image.Image]: | |
| if not codes: | |
| return [] | |
| shapes = {tuple(value.shape) for value in codes} | |
| if len(shapes) != 1: | |
| raise ValueError(f"same-grid decode batch has mixed code shapes: {shapes}") | |
| binary = torch.cat(codes, dim=0).to(self.device, torch.bfloat16) | |
| indices = self.decoder.bsq_encoder.bsq.codes_to_indexes(binary * 2 - 1) | |
| grid_h, grid_w = grid | |
| return self.decoder.decode( | |
| indices, | |
| ar_height=grid_h * 32, | |
| ar_width=grid_w * 32, | |
| upsampling_ratio=upsampling, | |
| num_inference_steps=steps, | |
| guidance_scale=guidance, | |
| generator=generators, | |
| ) | |
| class VisualMetricModels: | |
| """PIQ, DINOv2, and CLIP models with fixed image-only embeddings.""" | |
| def __init__(self, *, device: torch.device, dino_path, clip_path): | |
| import piq | |
| from transformers import ( | |
| AutoImageProcessor, | |
| AutoModel, | |
| CLIPVisionModelWithProjection, | |
| ) | |
| self.device = device | |
| self.piq = piq | |
| self.lpips = piq.LPIPS(reduction="none").eval().to(device) | |
| self.dists = piq.DISTS(reduction="none").eval().to(device) | |
| self.dino_processor = AutoImageProcessor.from_pretrained( | |
| dino_path, local_files_only=True, use_fast=False | |
| ) | |
| self.dino = AutoModel.from_pretrained( | |
| dino_path, | |
| local_files_only=True, | |
| dtype=torch.bfloat16, | |
| ).eval().to(device) | |
| self.clip_processor = AutoImageProcessor.from_pretrained( | |
| clip_path, local_files_only=True, use_fast=False | |
| ) | |
| self.clip = CLIPVisionModelWithProjection.from_pretrained( | |
| clip_path, | |
| local_files_only=True, | |
| dtype=torch.bfloat16, | |
| ).eval().to(device) | |
| def embeddings( | |
| self, images: list[Image.Image], *, batch_size: int | |
| ) -> tuple[torch.Tensor, torch.Tensor]: | |
| if batch_size <= 0: | |
| raise ValueError("semantic batch size must be positive") | |
| dino_values = [] | |
| clip_values = [] | |
| for start in range(0, len(images), batch_size): | |
| chunk = [image.convert("RGB") for image in images[start : start + batch_size]] | |
| dino_inputs = self.dino_processor(images=chunk, return_tensors="pt") | |
| dino_pixels = dino_inputs["pixel_values"].to( | |
| self.device, dtype=torch.bfloat16 | |
| ) | |
| dino_output = self.dino(pixel_values=dino_pixels) | |
| dino_embedding = dino_output.last_hidden_state[:, 0].float() | |
| dino_values.append(F.normalize(dino_embedding, dim=-1).cpu()) | |
| clip_inputs = self.clip_processor(images=chunk, return_tensors="pt") | |
| clip_pixels = clip_inputs["pixel_values"].to( | |
| self.device, dtype=torch.bfloat16 | |
| ) | |
| clip_embedding = self.clip(pixel_values=clip_pixels).image_embeds.float() | |
| clip_values.append(F.normalize(clip_embedding, dim=-1).cpu()) | |
| return torch.cat(dino_values), torch.cat(clip_values) | |
| def aligned_metrics( | |
| self, generated: Image.Image, reference: Image.Image | |
| ) -> dict: | |
| generated_tensor, reference_tensor = align_generated_to_reference( | |
| generated, reference | |
| ) | |
| generated_tensor = generated_tensor.to(self.device) | |
| reference_tensor = reference_tensor.to(self.device) | |
| minimum = min(generated_tensor.shape[-2:]) | |
| levels = min( | |
| len(_MS_SSIM_WEIGHTS), | |
| max(1, int(math.floor(math.log2(max((minimum - 1) / 10, 1)))) + 1), | |
| ) | |
| weights = generated_tensor.new_tensor(_MS_SSIM_WEIGHTS[:levels]) | |
| mse = F.mse_loss(generated_tensor, reference_tensor) | |
| psnr = -10 * torch.log10(mse.clamp_min(1e-12)) | |
| return { | |
| "mse": finite_number(mse), | |
| "psnr_db": finite_number(psnr), | |
| "ssim": finite_number( | |
| self.piq.ssim( | |
| generated_tensor, | |
| reference_tensor, | |
| data_range=1.0, | |
| reduction="mean", | |
| ) | |
| ), | |
| "ms_ssim": finite_number( | |
| self.piq.multi_scale_ssim( | |
| generated_tensor, | |
| reference_tensor, | |
| data_range=1.0, | |
| reduction="mean", | |
| scale_weights=weights, | |
| ) | |
| ), | |
| "ms_ssim_levels": levels, | |
| "lpips": finite_number(self.lpips(generated_tensor, reference_tensor).mean()), | |
| "dists": finite_number(self.dists(generated_tensor, reference_tensor).mean()), | |
| } | |
| def aligned_lpips( | |
| self, first: Image.Image, second: Image.Image | |
| ) -> float: | |
| first_tensor, second_tensor = align_generated_to_reference(first, second) | |
| return finite_number( | |
| self.lpips( | |
| first_tensor.to(self.device), second_tensor.to(self.device) | |
| ).mean() | |
| ) | |
| __all__ = ["RolloutRGBDecoder", "VisualMetricModels"] | |
Xet Storage Details
- Size:
- 6.12 kB
- Xet hash:
- ab645f097ab9c87a3a5b16580e636153cbd3c4ebca809ad046bbb46187cdb56a
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.