import time import random import numpy as np import torch from PIL import Image from modules import devices from modules.shared import opts, log from modules.upscaler import Upscaler, UpscalerData from modules.image import convert from modules.model_quant import do_post_load_quant MODELS_MAP = { "SeedVR2 3B": "seedvr2_ema_3b_fp16.safetensors", "SeedVR2 7B": "seedvr2_ema_7b_fp16.safetensors", "SeedVR2 7B Sharp": "seedvr2_ema_7b_sharp_fp16.safetensors", } class UpscalerSeedVR(Upscaler): def __init__(self, dirname=None): self.name = "SeedVR2" super().__init__() self.scalers = [ UpscalerData(name="SeedVR2 3B", path=None, upscaler=self, model=None, scale=1), UpscalerData(name="SeedVR2 7B", path=None, upscaler=self, model=None, scale=1), UpscalerData(name="SeedVR2 7B Sharp", path=None, upscaler=self, model=None, scale=1), ] self.model = None self.model_loaded = None self.tile_size = 1024 self.tile_overlap = 0.25 self.device = devices.device def load_model(self, path: str): model_name = MODELS_MAP.get(path, None) if (self.model is None) or (self.model_loaded != model_name): log.debug(f'Upscaler loading: name="{self.name}" model="{model_name}"') t0 = time.time() from modules.seedvr.src.core.model_manager import configure_runner from modules.seedvr.src.core import generation self.model = configure_runner( model_name=model_name, cache_dir=opts.hfcache_dir, device=devices.device, dtype=devices.dtype, ) self.model_loaded = model_name self.model.dit.device = devices.device self.model.dit.dtype = devices.dtype self.model.vae_encode = self.vae_encode self.model.vae_decode = self.vae_decode # Patch generation_loop's generation_step() with our wrapper; stash the original once # so reloads don't re-wrap the wrapper itself (infinite recursion). if not hasattr(generation, "generation_step_original"): generation.generation_step_original = generation.generation_step generation.generation_step = self.model_step self.model._internal_dict = { 'dit': self.model.dit, 'vae': self.model.vae, } t1 = time.time() self.model.dit.config = self.model.config.dit self.model.vae.tile_sample_min_size = self.tile_size self.model.vae.tile_latent_min_size = self.tile_size // 8 self.model.vae.tile_overlap_factor = self.tile_overlap self.model = do_post_load_quant(self.model, allow=True) log.info(f'Upscaler loaded: name="{self.name}" model="{model_name}" time={t1 - t0:.2f}') def vae_encode(self, samples): log.debug(f'Upscaler encode: samples={samples[0].shape if len(samples) > 0 else None} tile={self.model.vae.tile_sample_min_size} overlap={self.model.vae.tile_overlap_factor}') latents = [] if len(samples) == 0: return latents self.model.dit = self.model.dit.to(device="cpu") self.model.vae = self.model.vae.to(device=self.device) devices.torch_gc() from einops import rearrange from modules.seedvr.src.optimization import memory_manager memory_manager.clear_rope_cache(self.model) scale = self.model.config.vae.scaling_factor shift = self.model.config.vae.get("shifting_factor", 0.0) batches = [sample.unsqueeze(0) for sample in samples] for sample in batches: sample = sample.to(self.device, self.model.vae.dtype) sample = self.model.vae.preprocess(sample) latent = self.model.vae.encode(sample).latent latent = latent.unsqueeze(2) if latent.ndim == 4 else latent latent = rearrange(latent, "b c ... -> b ... c") latent = (latent - shift) * scale latents.append(latent) latents = [latent.squeeze(0) for latent in latents] self.model.vae = self.model.vae.to(device="cpu") devices.torch_gc() return latents def vae_decode(self, latents, target_dtype: torch.dtype = None): log.debug(f'Upscaler decode: latents={latents[0].shape if len(latents) > 0 else None} tile={self.model.vae.tile_latent_min_size} overlap={self.model.vae.tile_overlap_factor}') samples = [] if len(latents) == 0: return samples from einops import rearrange from modules.seedvr.src.optimization import memory_manager memory_manager.clear_rope_cache(self.model) self.model.dit = self.model.dit.to(device="cpu") self.model.vae = self.model.vae.to(device=self.device) devices.torch_gc() scale = self.model.config.vae.scaling_factor shift = self.model.config.vae.get("shifting_factor", 0.0) latents = [latent.unsqueeze(0) for latent in latents] with devices.inference_context(): for _i, latent in enumerate(latents): latent = latent.to(self.device, self.model.vae.dtype) latent = latent / scale + shift latent = rearrange(latent, "b ... c -> b c ...") latent = latent.squeeze(2) sample = self.model.vae.decode(latent).sample sample = self.model.vae.postprocess(sample) samples.append(sample) samples = [sample.squeeze(0) for sample in samples] self.model.vae = self.model.vae.to(device="cpu") devices.torch_gc() return samples def model_step(self, *args, **kwargs): from modules.seedvr.src.core import generation from modules.seedvr.src.optimization import memory_manager self.model.vae = self.model.vae.to(device="cpu") self.model.dit = self.model.dit.to(device=self.device) devices.torch_gc() log.debug(f'Upscaler inference: args={len(args)} kwargs={list(kwargs.keys())}') memory_manager.preinitialize_rope_cache(self.model) with devices.inference_context(): result = generation.generation_step_original(*args, **kwargs) self.model.dit = self.model.dit.to(device="cpu") devices.torch_gc() return result def do_upscale(self, img: Image.Image, selected_file, cfg_scale: float = 3.5, cfg_rescale: float = 0.0, steps: int = 1, seed: int = -1, scale: float | None = None, tile_size: int = 1024, tile_overlap: float = 0.25): self.load_model(selected_file) if self.model is None: return img from modules.seedvr.src.core import generation self.scale = self.scale if scale is None else scale self.tile_size = tile_size if tile_size is not None else self.tile_size self.tile_overlap = tile_overlap if tile_overlap is not None else self.tile_overlap self.model.vae.tile_sample_min_size = self.tile_size self.model.vae.tile_latent_min_size = self.tile_size // 8 self.model.vae.tile_overlap_factor = self.tile_overlap width = int(self.scale * img.width) // 8 * 8 image_tensor = np.array(img) image_tensor = torch.from_numpy(image_tensor).to(device=devices.device, dtype=devices.dtype).unsqueeze(0) / 255.0 random.seed() seed = int(random.randrange(4294967294)) if seed == -1 else int(seed) t0 = time.time() with devices.inference_context(): result_tensor = generation.generation_loop( runner=self.model, images=image_tensor, cfg_scale=cfg_scale, cfg_rescale=cfg_rescale, steps=steps, seed=seed, res_w=width, batch_size=1, temporal_overlap=0, device=devices.device, ) t1 = time.time() tiles = getattr(self.model.vae, "tiles", None) log.info(f'Upscaler: type="{self.name}" model="{selected_file}" scale={self.scale} cfg={cfg_scale} seed={seed} tiles={tiles} time={t1 - t0:.2f}') img = convert.to_pil(result_tensor.squeeze()) if opts.upscaler_unload: self.model.dit = None self.model.vae = None self.model.cache = None self.model = None log.debug(f'Upscaler unload: type="{self.name}" model="{selected_file}"') devices.torch_gc(force=True) return img