Skip to content

[Sana-4K] #10537

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 10 commits into from
Jan 14, 2025
+23 −4
Merged

[Sana-4K] #10537

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
8232a2f
[Sana 4K]
lawrence-cj Jan 8, 2025
e626e34
Merge branch 'main' into Sana-4K
lawrence-cj Jan 8, 2025
82edf53
Merge branch 'main' into Sana-4K
lawrence-cj Jan 12, 2025
b00ce80
[Sana-4K] fix SanaPAGPipeline
lawrence-cj Jan 12, 2025
4d18b53
add VAE automatically tiling function;
lawrence-cj Jan 12, 2025
8fe11f4
set clean_caption to False;
lawrence-cj Jan 12, 2025
f5dff32
Merge branch 'main' into Sana-4K
lawrence-cj Jan 13, 2025
45fadaf
Merge branch 'main' into Sana-4K
lawrence-cj Jan 14, 2025
0932a4e
add warnings for VAE OOM.
lawrence-cj Jan 14, 2025
422fd36
style
yiyixuxu Jan 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 12 additions & 3 deletions src/diffusers/pipelines/pag/pipeline_pag_sana.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@
ASPECT_RATIO_1024_BIN,
)
from ..pixart_alpha.pipeline_pixart_sigma import ASPECT_RATIO_2048_BIN
from ..sana.pipeline_sana import ASPECT_RATIO_4096_BIN
from .pag_utils import PAGMixin


Expand Down Expand Up @@ -639,7 +640,7 @@ def __call__(
negative_prompt_attention_mask: Optional[torch.Tensor] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
clean_caption: bool = True,
clean_caption: bool = False,
use_resolution_binning: bool = True,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
Expand Down Expand Up @@ -755,7 +756,9 @@ def __call__(
callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs

if use_resolution_binning:
if self.transformer.config.sample_size == 64:
if self.transformer.config.sample_size == 128:
aspect_ratio_bin = ASPECT_RATIO_4096_BIN
elif self.transformer.config.sample_size == 64:
aspect_ratio_bin = ASPECT_RATIO_2048_BIN
elif self.transformer.config.sample_size == 32:
aspect_ratio_bin = ASPECT_RATIO_1024_BIN
Expand Down Expand Up @@ -912,7 +915,13 @@ def __call__(
image = latents
else:
latents = latents.to(self.vae.dtype)
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
try:
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
except torch.cuda.OutOfMemoryError as e:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok here, if you feel strongly about this, we prefer to throw out a warning with recommendations and user can explicitly enable the optimization

print("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")
self.vae.enable_tiling(tile_sample_min_width=1024, tile_sample_min_height=1024)
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
self.vae.disable_tiling()
if use_resolution_binning:
image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height)

Expand Down
8 changes: 7 additions & 1 deletion src/diffusers/pipelines/sana/pipeline_sana.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -953,7 +953,13 @@ def __call__(
image = latents
else:
latents = latents.to(self.vae.dtype)
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
try:
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
except torch.cuda.OutOfMemoryError as e:
print("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")
self.vae.enable_tiling(tile_sample_min_width=1024, tile_sample_min_height=1024)
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
self.vae.disable_tiling()
if use_resolution_binning:
image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height)

Expand Down
Loading