AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-27 17:22:53 +03:00
Code Activity

[Paint by example] Fix cpu offload for paint by example (#2062)

Browse Source 
* [Paint by example] Fix paint by example

* fix more

* final fix
This commit is contained in:
Patrick von Platen
2023-01-24 11:05:50 +02:00
committed by GitHub
parent a66f2baeb7
commit 7d8b4f7f8e
2 changed files with 8 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/diffusers/pipelines/paint_by_example/image_encoder.py
View File

@@ -36,12 +36,15 @@ class PaintByExampleImageEncoder(CLIPPreTrainedModel):
# uncondition for scaling
self.uncond_vector = nn.Parameter(torch.randn((1, 1, self.proj_size)))
def forward(self, pixel_values):
def forward(self, pixel_values, return_uncond_vector=False):
clip_output = self.model(pixel_values=pixel_values)
latent_states = clip_output.pooler_output
latent_states = self.mapper(latent_states[:, None])
latent_states = self.final_layer_norm(latent_states)
latent_states = self.proj_out(latent_states)
if return_uncond_vector:
return latent_states, self.uncond_vector
return latent_states

12
src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py
View File

@@ -201,14 +201,11 @@ class PaintByExamplePipeline(DiffusionPipeline):
device = torch.device(f"cuda:{gpu_id}")
for cpu_offloaded_model in [self.unet, self.vae]:
if cpu_offloaded_model is not None:
cpu_offload(cpu_offloaded_model, device)
for cpu_offloaded_model in [self.unet, self.vae, self.image_encoder]:
cpu_offload(cpu_offloaded_model, execution_device=device)
if self.safety_checker is not None:
# TODO(Patrick) - there is currently a bug with cpu offload of nn.Parameter in accelerate
# fix by only offloading self.safety_checker for now
cpu_offload(self.safety_checker.vision_model, device)
cpu_offload(self.safety_checker, execution_device=device, offload_buffers=True)
@property
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
@@ -367,7 +364,7 @@ class PaintByExamplePipeline(DiffusionPipeline):
image = self.feature_extractor(images=image, return_tensors="pt").pixel_values
image = image.to(device=device, dtype=dtype)
image_embeddings = self.image_encoder(image)
image_embeddings, uncond_embeddings = self.image_encoder(image, return_uncond_vector=True)
# duplicate image embeddings for each generation per prompt, using mps friendly method
bs_embed, seq_len, _ = image_embeddings.shape
@@ -375,7 +372,6 @@ class PaintByExamplePipeline(DiffusionPipeline):
image_embeddings = image_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)
if do_classifier_free_guidance:
uncond_embeddings = self.image_encoder.uncond_vector
uncond_embeddings = uncond_embeddings.repeat(1, image_embeddings.shape[0], 1)
uncond_embeddings = uncond_embeddings.view(bs_embed * num_images_per_prompt, 1, -1)