Spaces:
Running on Zero
Running on Zero
| """Texture latent codec for SEFI-T2I.""" | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| from torch import Tensor | |
| class TextureLatentCodec(nn.Module): | |
| """Encode/decode and normalize texture latents for SEFI.""" | |
| def __init__( | |
| self, | |
| texture_vae: nn.Module, | |
| texture_vae_name: str, | |
| ): | |
| super().__init__() | |
| self.texture_vae = texture_vae | |
| self.texture_vae_name = str(texture_vae_name) | |
| self._use_flux2_bn = self.texture_vae_name == "flux2" | |
| config = getattr(texture_vae, "config", None) | |
| latent_channels = getattr(config, "latent_channels", None) | |
| if latent_channels is None: | |
| raise ValueError( | |
| "Texture VAE config must provide latent_channels for channel derivation." | |
| ) | |
| self.latent_channels = int(latent_channels) | |
| self.texture_channels = int(self.latent_channels * 4) | |
| if self._use_flux2_bn: | |
| if not hasattr(texture_vae, "bn"): | |
| raise ValueError( | |
| f"Texture VAE '{self.texture_vae_name}' requires bn stats but no bn module found." | |
| ) | |
| eps = float(getattr(config, "batch_norm_eps", 1e-6)) | |
| bn_mean = texture_vae.bn.running_mean.view(1, -1, 1, 1).float() | |
| bn_std = torch.sqrt(texture_vae.bn.running_var.view(1, -1, 1, 1).float() + eps) | |
| self.register_buffer("texture_bn_mean", bn_mean, persistent=False) | |
| self.register_buffer("texture_bn_std", bn_std, persistent=False) | |
| self.scaling_factor = None | |
| self.shift_factor = None | |
| else: | |
| scaling_factor = float(getattr(config, "scaling_factor", 1.0)) | |
| shift_factor = float(getattr(config, "shift_factor", 0.0) or 0.0) | |
| if scaling_factor <= 0: | |
| raise ValueError( | |
| f"Invalid scaling_factor={scaling_factor} for texture VAE {self.texture_vae_name}." | |
| ) | |
| self.scaling_factor = scaling_factor | |
| self.shift_factor = shift_factor | |
| def vae_dtype(self) -> torch.dtype: | |
| return next(self.texture_vae.parameters()).dtype | |
| def _encode_raw(self, images: Tensor) -> Tensor: | |
| return self.texture_vae.encode(images.to(dtype=self.vae_dtype)).latent_dist.mode() | |
| def _normalize_raw(self, raw_latents: Tensor) -> Tensor: | |
| return (raw_latents - self.shift_factor) * self.scaling_factor | |
| def _denormalize_raw(self, normed_latents: Tensor) -> Tensor: | |
| return normed_latents / self.scaling_factor + self.shift_factor | |
| def _normalize_patchified(self, patchified_latents: Tensor) -> Tensor: | |
| bn_mean = self.texture_bn_mean.to(patchified_latents.device, patchified_latents.dtype) | |
| bn_std = self.texture_bn_std.to(patchified_latents.device, patchified_latents.dtype) | |
| return (patchified_latents - bn_mean) / bn_std | |
| def _denormalize_patchified(self, patchified_latents: Tensor) -> Tensor: | |
| bn_mean = self.texture_bn_mean.to(patchified_latents.device, patchified_latents.dtype) | |
| bn_std = self.texture_bn_std.to(patchified_latents.device, patchified_latents.dtype) | |
| return patchified_latents * bn_std + bn_mean | |
| def encode_texture(self, images: Tensor, pipeline_cls) -> Tensor: | |
| raw_latents = self._encode_raw(images) | |
| if self._use_flux2_bn: | |
| patchified = pipeline_cls._patchify_latents(raw_latents) | |
| patchified = self._normalize_patchified(patchified) | |
| else: | |
| normed_raw = self._normalize_raw(raw_latents) | |
| patchified = pipeline_cls._patchify_latents(normed_raw) | |
| if patchified.shape[1] != self.texture_channels: | |
| raise ValueError( | |
| f"Texture channels mismatch: derived={self.texture_channels}, got={patchified.shape[1]}." | |
| ) | |
| return patchified | |
| def decode_texture(self, texture_latents: Tensor, pipeline_cls) -> Tensor: | |
| if self._use_flux2_bn: | |
| patchified = self._denormalize_patchified(texture_latents) | |
| raw_latents = pipeline_cls._unpatchify_latents(patchified) | |
| else: | |
| raw_normed = pipeline_cls._unpatchify_latents(texture_latents) | |
| raw_latents = self._denormalize_raw(raw_normed) | |
| return self.texture_vae.decode(raw_latents, return_dict=False)[0] | |