下沉vae tiled相关参数到base (#50)

Author: tenderness-git

diffsynth_engine/pipelines/base.py

@@ -25,11 +25,14 @@ def convert(self, lora_state_dict: Dict[str, torch.Tensor]) -> Dict[str, Dict[st
 class BasePipeline:
     lora_converter = LoRAStateDictConverter()
 
-    def __init__(self, device="cuda:0", dtype=torch.float16):
+    def __init__(self, vae_tiled, vae_tile_size, vae_tile_stride, device="cuda:0", dtype=torch.float16):
         super().__init__()
         self.device = device
         self.dtype = dtype
         self.offload_mode = None
+        self.vae_tiled = vae_tiled
+        self.vae_tile_size = vae_tile_size
+        self.vae_tile_stride = vae_tile_stride
         self.model_names = []
 
     @classmethod
@@ -140,13 +143,17 @@ def generate_noise(shape, seed=None, device="cpu", dtype=torch.float16):
         noise = torch.randn(shape, generator=generator, device=device, dtype=dtype)
         return noise
 
-    def encode_image(self, image: torch.Tensor, tiled=False, tile_size=64, tile_stride=32) -> torch.Tensor:
-        latents = self.vae_encoder(image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)
+    def encode_image(self, image: torch.Tensor) -> torch.Tensor:
+        latents = self.vae_encoder(
+            image, tiled=self.vae_tiled, tile_size=self.vae_tile_size, tile_stride=self.vae_tile_stride
+        )
         return latents
 
-    def decode_image(self, latent: torch.Tensor, tiled=False, tile_size=64, tile_stride=32) -> torch.Tensor:
+    def decode_image(self, latent: torch.Tensor) -> torch.Tensor:
         vae_dtype = self.vae_decoder.conv_in.weight.dtype
-        image = self.vae_decoder(latent.to(vae_dtype), tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)
+        image = self.vae_decoder(
+            latent.to(vae_dtype), tiled=self.vae_tiled, tile_size=self.vae_tile_size, tile_stride=self.vae_tile_stride
+        )
         return image
 
     def prepare_latents(

diffsynth_engine/pipelines/flux_image.py

@@ -225,10 +225,19 @@ def __init__(
         vae_encoder: FluxVAEEncoder,
         use_cfg: bool = False,
         batch_cfg: bool = False,
+        vae_tiled: bool = False,
+        vae_tile_size: int = 256,
+        vae_tile_stride: int = 256,
         device: str = "cuda:0",
         dtype: torch.dtype = torch.bfloat16,
     ):
-        super().__init__(device=device, dtype=dtype)
+        super().__init__(
+            vae_tiled=vae_tiled,
+            vae_tile_size=vae_tile_size,
+            vae_tile_stride=vae_tile_stride,
+            device=device,
+            dtype=dtype,
+        )
         self.noise_scheduler = RecifitedFlowScheduler(shift=3.0, use_dynamic_shifting=True)
         self.sampler = FlowMatchEulerSampler()
         # models
@@ -474,9 +483,6 @@ def prepare_latents(
         denoising_strength: float,
         num_inference_steps: int,
         mu: float,
-        tiled: bool = False,
-        tile_size: int = 128,
-        tile_stride: int = 64,
     ):
         # Prepare scheduler
         if input_image is not None:
@@ -491,7 +497,7 @@ def prepare_latents(
             self.load_models_to_device(["vae_encoder"])
             noise = latents
             image = self.preprocess_image(input_image).to(device=self.device, dtype=self.dtype)
-            latents = self.encode_image(image, tiled, tile_size, tile_stride)
+            latents = self.encode_image(image)
             init_latents = latents.clone()
             latents = self.sampler.add_noise(latents, noise, sigma_start)
         else:
@@ -506,15 +512,15 @@ def prepare_masked_latent(self, image: Image.Image, mask: Image.Image | None, he
         if mask is None:
             image = image.resize((width, height))
             image = self.preprocess_image(image).to(device=self.device, dtype=self.dtype)
-            latent = self.encode_image(image, tiled=False)
+            latent = self.encode_image(image)
         else:
             image = image.resize((width, height))
             mask = mask.resize((width, height))
             image = self.preprocess_image(image).to(device=self.device, dtype=self.dtype)
             mask = self.preprocess_mask(mask).to(device=self.device, dtype=self.dtype)
             masked_image = image.clone()
             masked_image[(mask > 0.5).repeat(1, 3, 1, 1)] = -1
-            latent = self.encode_image(masked_image, tiled=False)
+            latent = self.encode_image(masked_image)
             mask = torch.nn.functional.interpolate(mask, size=(latent.shape[2], latent.shape[3]))
             mask = 1 - mask
             latent = torch.cat([latent, mask], dim=1)
@@ -585,9 +591,6 @@ def __call__(
         height: int = 1024,
         width: int = 1024,
         num_inference_steps: int = 30,
-        tiled: bool = False,
-        tile_size: int = 128,
-        tile_stride: int = 64,
         seed: int | None = None,
         controlnet_params: List[ControlNetParams] | ControlNetParams = [],
         progress_callback: Optional[Callable] = None,  # def progress_callback(current, total, status)
@@ -605,7 +608,7 @@ def __call__(
         image_seq_len = math.ceil(height // 16) * math.ceil(width // 16)
         mu = calculate_shift(image_seq_len)
         init_latents, latents, sigmas, timesteps = self.prepare_latents(
-            noise, input_image, denoising_strength, num_inference_steps, mu, tiled, tile_size, tile_stride
+            noise, input_image, denoising_strength, num_inference_steps, mu
         )
         # Initialize sampler
         self.sampler.initialize(init_latents=init_latents, timesteps=timesteps, sigmas=sigmas)
@@ -649,7 +652,7 @@ def __call__(
                 progress_callback(i, len(timesteps), "DENOISING")
         # Decode image
         self.load_models_to_device(["vae_decoder"])
-        vae_output = self.decode_image(latents, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)
+        vae_output = self.decode_image(latents)
         image = self.vae_output_to_image(vae_output)
         # Offload all models
         self.load_models_to_device([])

diffsynth_engine/pipelines/sd_image.py

@@ -155,10 +155,19 @@ def __init__(
         vae_decoder: SDVAEDecoder,
         vae_encoder: SDVAEEncoder,
         batch_cfg: bool = True,
+        vae_tiled: bool = False,
+        vae_tile_size: int = 256,
+        vae_tile_stride: int = 256,
         device: str = "cuda",
         dtype: torch.dtype = torch.float16,
     ):
-        super().__init__(device=device, dtype=dtype)
+        super().__init__(
+            vae_tiled=vae_tiled,
+            vae_tile_size=vae_tile_size,
+            vae_tile_stride=vae_tile_stride,
+            device=device,
+            dtype=dtype,
+        )
         self.noise_scheduler = ScaledLinearScheduler()
         self.sampler = EulerSampler()
         # models
@@ -310,9 +319,6 @@ def __call__(
         height: int = 1024,
         width: int = 1024,
         num_inference_steps: int = 20,
-        tiled: bool = False,
-        tile_size: int = 64,
-        tile_stride: int = 32,
         seed: int | None = None,
         progress_callback: Optional[Callable] = None,  # def progress_callback(current, total, status)
     ):
@@ -322,7 +328,7 @@ def __call__(
         noise = self.generate_noise((1, 4, height // 8, width // 8), seed=seed, device=self.device, dtype=self.dtype)
 
         init_latents, latents, sigmas, timesteps = self.prepare_latents(
-            noise, input_image, denoising_strength, num_inference_steps, tiled, tile_size, tile_stride
+            noise, input_image, denoising_strength, num_inference_steps
         )
         mask, overlay_image = None, None
         if mask_image is not None:
@@ -359,7 +365,7 @@ def __call__(
             latents = latents * mask + init_latents * (1 - mask)
         # Decode image
         self.load_models_to_device(["vae_decoder"])
-        vae_output = self.decode_image(latents, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)
+        vae_output = self.decode_image(latents)
         image = self.vae_output_to_image(vae_output)
         # Paste Overlay Image
         if mask_image is not None:

diffsynth_engine/pipelines/sdxl_image.py

@@ -124,10 +124,19 @@ def __init__(
         vae_decoder: SDXLVAEDecoder,
         vae_encoder: SDXLVAEEncoder,
         batch_cfg: bool = True,
+        vae_tiled: bool = False,
+        vae_tile_size: int = 256,
+        vae_tile_stride: int = 256,
         device: str = "cuda",
         dtype: torch.dtype = torch.float16,
     ):
-        super().__init__(device=device, dtype=dtype)
+        super().__init__(
+            vae_tiled=vae_tiled,
+            vae_tile_size=vae_tile_size,
+            vae_tile_stride=vae_tile_stride,
+            device=device,
+            dtype=dtype,
+        )
         self.noise_scheduler = ScaledLinearScheduler()
         self.sampler = EulerSampler()
         # models
@@ -342,9 +351,6 @@ def __call__(
         height: int = 1024,
         width: int = 1024,
         num_inference_steps: int = 20,
-        tiled: bool = False,
-        tile_size: int = 64,
-        tile_stride: int = 32,
         seed: int | None = None,
         progress_callback: Optional[Callable] = None,  # def progress_callback(current, total, status)
     ):
@@ -354,7 +360,7 @@ def __call__(
         noise = self.generate_noise((1, 4, height // 8, width // 8), seed=seed, device=self.device, dtype=self.dtype)
 
         init_latents, latents, sigmas, timesteps = self.prepare_latents(
-            noise, input_image, denoising_strength, num_inference_steps, tiled, tile_size, tile_stride
+            noise, input_image, denoising_strength, num_inference_steps
         )
         mask, overlay_image = None, None
         if mask_image is not None:
@@ -402,7 +408,7 @@ def __call__(
             latents = latents * mask + init_latents * (1 - mask)
         # Decode image
         self.load_models_to_device(["vae_decoder"])
-        vae_output = self.decode_image(latents, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)
+        vae_output = self.decode_image(latents)
         image = self.vae_output_to_image(vae_output)
 
         if mask_image is not None:

diffsynth_engine/pipelines/wan_video.py

@@ -128,10 +128,19 @@ def __init__(
         vae: WanVideoVAE,
         image_encoder: WanImageEncoder,
         batch_cfg: bool = False,
+        vae_tiled: bool = True,
+        vae_tile_size: Tuple[int, int] = (34, 34),
+        vae_tile_stride: Tuple[int, int] = (18, 16),
         device="cuda",
         dtype=torch.bfloat16,
     ):
-        super().__init__(device=device, dtype=dtype)
+        super().__init__(
+            vae_tiled=vae_tiled,
+            vae_tile_size=vae_tile_size,
+            vae_tile_stride=vae_tile_stride,
+            device=device,
+            dtype=dtype,
+        )
         self.noise_scheduler = RecifitedFlowScheduler(shift=5.0, sigma_min=0.001, sigma_max=0.999)
         self.sampler = FlowMatchEulerSampler()
         self.tokenizer = tokenizer
@@ -202,22 +211,26 @@ def tensor2video(self, frames):
         frames = [Image.fromarray(frame) for frame in frames]
         return frames
 
-    def encode_video(self, videos: torch.Tensor, tiled=True, tile_size=(34, 34), tile_stride=(18, 16)):
+    def encode_video(self, videos: torch.Tensor):
         videos = videos.to(dtype=self.config.vae_dtype, device=self.device)
-        latents = self.vae.encode(videos, device=self.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)
+        latents = self.vae.encode(
+            videos,
+            device=self.device,
+            tiled=self.vae_tiled,
+            tile_size=self.vae_tile_size,
+            tile_stride=self.vae_tile_stride,
+        )
         latents = latents.to(dtype=self.config.dit_dtype, device=self.device)
         return latents
 
-    def decode_video(
-        self, latents, tiled=True, tile_size=(34, 34), tile_stride=(18, 16), progress_callback=None
-    ) -> List[torch.Tensor]:
+    def decode_video(self, latents, progress_callback=None) -> List[torch.Tensor]:
         latents = latents.to(dtype=self.config.vae_dtype, device=self.device)
         videos = self.vae.decode(
             latents,
             device=self.device,
-            tiled=tiled,
-            tile_size=tile_size,
-            tile_stride=tile_stride,
+            tiled=self.vae_tiled,
+            tile_size=self.vae_tile_size,
+            tile_stride=self.vae_tile_stride,
             progress_callback=progress_callback,
         )
         videos = [video.to(dtype=self.config.dit_dtype, device=self.device) for video in videos]
@@ -297,9 +310,6 @@ def prepare_latents(
         input_video,
         denoising_strength,
         num_inference_steps,
-        tiled=True,
-        tile_size=(34, 34),
-        tile_stride=(18, 16),
     ):
         if input_video is not None:
             total_steps = num_inference_steps
@@ -311,9 +321,7 @@ def prepare_latents(
             noise = latents
             input_video = self.preprocess_images(input_video)
             input_video = torch.stack(input_video, dim=2)
-            latents = self.encode_video(input_video, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride).to(
-                dtype=latents.dtype, device=latents.device
-            )
+            latents = self.encode_video(input_video).to(dtype=latents.dtype, device=latents.device)
             init_latents = latents.clone()
             latents = self.sampler.add_noise(latents, noise, sigma_start)
         else:
@@ -336,9 +344,6 @@ def __call__(
         num_frames=81,
         cfg_scale=5.0,
         num_inference_steps=50,
-        tiled=True,
-        tile_size=(34, 34),
-        tile_stride=(18, 16),
         progress_callback: Optional[Callable] = None,  # def progress_callback(current, total, status)
     ):
         assert height % 16 == 0 and width % 16 == 0, "height and width must be divisible by 16"
@@ -353,9 +358,6 @@ def __call__(
             input_video,
             denoising_strength,
             num_inference_steps,
-            tiled=tiled,
-            tile_size=tile_size,
-            tile_stride=tile_stride,
         )
         self.sampler.initialize(init_latents=init_latents, timesteps=timesteps, sigmas=sigmas)
         # Encode prompts
@@ -392,9 +394,7 @@ def __call__(
 
         # Decode
         self.load_models_to_device(["vae"])
-        frames = self.decode_video(
-            latents, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride, progress_callback=progress_callback
-        )
+        frames = self.decode_video(latents, progress_callback=progress_callback)
         frames = self.tensor2video(frames[0])
         return frames