Spaces:
Running on Zero
Running on Zero
File size: 2,527 Bytes
e2ddf3f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | # PID (PixelDiT SR) model — inference subset.
#
# At inference the only thing this class adds on top of PixelDiTModel is the
# frozen VAE (`vae_encoder`) used by `encode_lq_latent`. The training-time
# degradation pipeline, LoRA injection, LPIPS loss, and training/validation
# steps have all been removed.
from __future__ import annotations
import logging
from typing import Any
import attrs
import torch
from torch import Tensor
from pid._ext.imaginaire.lazy_config import instantiate as lazy_instantiate
from pid._ext.imaginaire.utils import misc
from pid._src.models.pixeldit_model import PixelDiTModel, PixelDiTModelConfig
logger = logging.getLogger(__name__)
@attrs.define(slots=False)
class PidModelConfig(PixelDiTModelConfig):
# "image" = LQ image only, "latent" = LQ latent only, "image_latent" = both.
lq_condition_type: str = "latent"
# Frozen VAE config for encoding LQ images to latent.
tokenizer: Any = None
# VAE latent channels (must match tokenizer.latent_ch).
state_ch: int = 16
# Fixed prompt override (training convenience kept here so checkpoints that set
# use_fixed_prompt=True still load).
use_fixed_prompt: bool = False
fixed_positive_prompt: str = ""
class PidModel(PixelDiTModel):
"""PID (PixelDiT SR) inference model (frozen VAE + LQ-conditioned student)."""
def __init__(self, config: PidModelConfig):
super().__init__(config)
if config.tokenizer is not None:
with misc.timer("PidModel: load_vae"):
from pid._src.tokenizers.base_vae import BaseVAE
self.vae_encoder: BaseVAE = lazy_instantiate(config.tokenizer)
if config.state_ch > 0:
assert self.vae_encoder.latent_ch == config.state_ch, (
f"latent_ch {self.vae_encoder.latent_ch} != state_ch {config.state_ch}"
)
else:
self.vae_encoder = None
logger.warning("No VAE configured — LQ latent encoding disabled.")
@torch.no_grad()
def encode_lq_latent(self, lq_image: Tensor) -> Tensor:
"""Encode an LQ image through the frozen VAE.
Args:
lq_image: [B, C, H_lq, W_lq] in [-1, 1].
Returns:
LQ latent [B, z_dim, zH, zW].
"""
if lq_image.ndim == 4:
lq_image = lq_image.unsqueeze(2)
latent = self.vae_encoder.encode(lq_image)
if latent.ndim == 5:
latent = latent[:, :, 0, :, :]
return latent
|