seedvt2

Signed-off-by: Vladimir Mandic <mandic00@live.com>

modules/postprocess/seedvr_model.py

@@ -1,9 +1,9 @@
 import time
+import random
 import numpy as np
 import torch
 from PIL import Image
 from torchvision.transforms import ToPILImage
-from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn
 from modules import devices
 from modules.shared import opts, log
 from modules.upscaler import Upscaler, UpscalerData
@@ -19,7 +19,7 @@ to_pil = ToPILImage()
 
 class UpscalerSeedVR(Upscaler):
     def __init__(self, dirname=None):
-        self.name = "SeedVR"
+        self.name = "SeedVR2"
         super().__init__()
         self.scalers = [
             UpscalerData(name="SeedVR2 3B", path=None, upscaler=self, model=None, scale=1),
@@ -32,45 +32,140 @@ class UpscalerSeedVR(Upscaler):
     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 load: name="{self.name}" model="{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
+            self.model.model_step = 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 = 1024
+            self.model.vae.tile_latent_min_size = 128
+            # from modules.model_quant import do_post_load_quant
+            # from modules.sd_offload import set_diffuser_offload
+            # self.model = do_post_load_quant(self.model, allow=True)
+            # set_diffuser_offload(self.model)
+            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.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 = self.model.model_step(*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):
-        devices.torch_gc()
         self.load_model(selected_file)
         if self.model is None:
             return img
 
-        from modules.seedvr.src.core.generation import generation_loop
+        from modules.seedvr.src.core import generation
 
         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))
+
         t0 = time.time()
-        result_tensor = generation_loop(
-            runner=self.model,
-            images=image_tensor,
-            cfg_scale=1.0,
-            seed=42,
-            res_w=width,
-            batch_size=1,
-            temporal_overlap=0,
-            device=devices.device,
-        )
+        with devices.inference_context():
+            result_tensor = generation.generation_loop(
+                runner=self.model,
+                images=image_tensor,
+                cfg_scale=opts.seedvt_cfg_scale,
+                seed=seed,
+                res_w=width,
+                batch_size=1,
+                temporal_overlap=0,
+                device=devices.device,
+            )
         t1 = time.time()
-        log.info(f'Upscaler: type="{self.name}" model="{selected_file}" scale={self.scale} time={t1 - t0:.2f}')
+        log.info(f'Upscaler: type="{self.name}" model="{selected_file}" scale={self.scale} cfg={opts.seedvt_cfg_scale} seed={seed} time={t1 - t0:.2f}')
         img = to_pil(result_tensor.squeeze().permute((2, 0, 1)))
-        devices.torch_gc()
 
         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)
+        devices.torch_gc(force=True)
         return img