github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Allow for vae tiling onle without low vram on wan models with a model kwarg

Browse Source
This commit is contained in:
Jaret Burkett
2026-07-10 09:42:52 -06:00
parent fe82487187
commit 886c2aec57
3 changed files with 27 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
extensions_built_in/diffusion_models/wan22/wan22_14b_model.py
View File

@@ -555,6 +555,11 @@ class Wan2214bModel(Wan21):
): ):
# reactivate progress bar since this is slooooow # reactivate progress bar since this is slooooow
pipeline.set_progress_bar_config(disable=False) pipeline.set_progress_bar_config(disable=False)
if self.use_vae_tiling:
# set vae to tile decode
pipeline.vae.enable_tiling()
# todo, figure out how to do video # todo, figure out how to do video
output = pipeline( output = pipeline(
prompt_embeds=conditional_embeds.text_embeds.to( prompt_embeds=conditional_embeds.text_embeds.to(
@@ -573,6 +578,10 @@ class Wan2214bModel(Wan21):
**extra **extra
)[0] )[0]
if self.use_vae_tiling:
# restore no tiling
pipeline.vae.disable_tiling()
# shape = [1, frames, channels, height, width] # shape = [1, frames, channels, height, width]
batch_item = output[0] # list of pil images batch_item = output[0] # list of pil images
if gen_config.num_frames > 1: if gen_config.num_frames > 1:

4
extensions_built_in/diffusion_models/wan22/wan22_5b_model.py
View File

@@ -210,7 +210,7 @@ class Wan225bModel(Wan21):
latent_model_input=latents, first_frame=first_frame_n1p1, vae=self.vae latent_model_input=latents, first_frame=first_frame_n1p1, vae=self.vae
) )
if self.model_config.low_vram: if self.use_vae_tiling:
# set vae to tile decode # set vae to tile decode
pipeline.vae.enable_tiling() pipeline.vae.enable_tiling()
@@ -234,7 +234,7 @@ class Wan225bModel(Wan21):
**extra, **extra,
)[0] )[0]
if self.model_config.low_vram: if self.use_vae_tiling:
# restore no tiling # restore no tiling
pipeline.vae.disable_tiling() pipeline.vae.disable_tiling()

16
toolkit/models/wan21/wan21.py
View File

@@ -520,6 +520,13 @@ class Wan21(BaseModel):
return pipeline return pipeline
@property
def use_vae_tiling(self):
# tile the vae decode when sampling if in low vram or explicitly enabled
return self.model_config.low_vram or self.model_config.model_kwargs.get(
"vae_tiling", False
)
def generate_single_image( def generate_single_image(
self, self,
pipeline: WanPipeline, pipeline: WanPipeline,
@@ -532,6 +539,11 @@ class Wan21(BaseModel):
# reactivate progress bar since this is slooooow # reactivate progress bar since this is slooooow
pipeline.set_progress_bar_config(disable=False) pipeline.set_progress_bar_config(disable=False)
pipeline = pipeline.to(self.device_torch) pipeline = pipeline.to(self.device_torch)
if self.use_vae_tiling:
# set vae to tile decode
pipeline.vae.enable_tiling()
# todo, figure out how to do video # todo, figure out how to do video
output = pipeline( output = pipeline(
prompt_embeds=conditional_embeds.text_embeds.to( prompt_embeds=conditional_embeds.text_embeds.to(
@@ -550,6 +562,10 @@ class Wan21(BaseModel):
**extra **extra
)[0] )[0]
if self.use_vae_tiling:
# restore no tiling
pipeline.vae.disable_tiling()
# shape = [1, frames, channels, height, width] # shape = [1, frames, channels, height, width]
batch_item = output[0] # list of pil images batch_item = output[0] # list of pil images
if gen_config.num_frames > 1: if gen_config.num_frames > 1: