From df02d15562c68a3c47a246085b88b43bad9dc8f1 Mon Sep 17 00:00:00 2001
From: Johannes Stelzer <jsdmail@gmail.com>
Date: Mon, 8 Jan 2024 20:31:28 +0100
Subject: [PATCH] fixed generator for prepare_extra_step_kwargs
---
diffusers_holder.py | 149 +++++---------------------------------------
1 file changed, 16 insertions(+), 133 deletions(-)
diff --git a/diffusers_holder.py b/diffusers_holder.py
index b0d381b..a65cdfe 100644
--- a/diffusers_holder.py
+++ b/diffusers_holder.py
@@ -114,8 +114,6 @@ class DiffusersHolder():
def get_noise(self, seed=420):
- generator = torch.Generator(device=self.device).manual_seed(int(seed))
-
latents = self.pipe.prepare_latents(
1,
self.pipe.unet.config.in_channels,
@@ -123,7 +121,7 @@ class DiffusersHolder():
self.width_img,
torch.float16,
self.pipe._execution_device,
- generator,
+ torch.Generator(device=self.device).manual_seed(int(seed)),
None,
)
@@ -448,6 +446,7 @@ class DiffusersHolder():
list_latents_mixing=None,
mixing_coeffs=0.0,
return_image: Optional[bool] = False,
+ seed=420,
**kwargs,
):
@@ -478,6 +477,7 @@ class DiffusersHolder():
clip_skip = None
callback_on_step_end = None
callback_on_step_end_tensor_inputs = ["latents"]
+
# 0. Default height and width to unet
height = self.pipe.default_sample_size * self.pipe.vae_scale_factor
@@ -488,8 +488,6 @@ class DiffusersHolder():
target_size = (height, width)
# 1. (skipped) Check inputs. Raise error if not correct
-
-
self.pipe._guidance_scale = guidance_scale
self.pipe._guidance_rescale = guidance_rescale
self.pipe._clip_skip = clip_skip
@@ -513,7 +511,8 @@ class DiffusersHolder():
list_latents_out = []
# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
- extra_step_kwargs = self.pipe.prepare_extra_step_kwargs(generator, eta)
+
+ extra_step_kwargs = self.pipe.prepare_extra_step_kwargs(torch.Generator(device=self.device).manual_seed(int(0)), eta)
# 7. Prepare added time ids & embeddings
add_text_embeds = pooled_prompt_embeds
@@ -627,120 +626,7 @@ class DiffusersHolder():
return self.latent2image(latents)
else:
return list_latents_out
-
-
- @torch.no_grad()
- def run_diffusion_sd_xl_old(
- self,
- text_embeddings: list,
- latents_start: torch.FloatTensor,
- idx_start: int = 0,
- list_latents_mixing=None,
- mixing_coeffs=0.0,
- return_image: Optional[bool] = False,
- **kwargs,
- ):
- # 0. Default height and width to unet
- original_size = (self.width_img, self.height_img)
- crops_coords_top_left = (0, 0)
- target_size = original_size
- batch_size = 1
- eta = 0.0
- num_images_per_prompt = 1
- cross_attention_kwargs = None
- generator = torch.Generator(device=self.device) # dummy generator
- do_classifier_free_guidance = self.guidance_scale > 1.0
-
- # 1. Check inputs. Raise error if not correct & 2. Define call parameters
- list_mixing_coeffs = self.prepare_mixing(mixing_coeffs, list_latents_mixing)
-
- # 3. Encode input prompt (already encoded outside bc of mixing, just split here)
- prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds = text_embeddings
-
- # 4. Prepare timesteps
- self.pipe.scheduler.set_timesteps(self.num_inference_steps, device=self.device)
- timesteps = self.pipe.scheduler.timesteps
-
- # 5. Prepare latent variables
- latents = latents_start.clone()
- list_latents_out = []
-
- # 6. Prepare extra step kwargs. usedummy generator
- extra_step_kwargs = self.pipe.prepare_extra_step_kwargs(generator, eta) # dummy
-
- # 7. Prepare added time ids & embeddings
- add_text_embeds = pooled_prompt_embeds
- if self.pipe.text_encoder_2 is None:
- text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
- else:
- text_encoder_projection_dim = self.pipe.text_encoder_2.config.projection_dim
-
- add_time_ids = self.pipe._get_add_time_ids(
- original_size,
- crops_coords_top_left,
- target_size,
- dtype=prompt_embeds.dtype,
- text_encoder_projection_dim=text_encoder_projection_dim,
- )
-
- negative_add_time_ids = add_time_ids
-
- prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
- add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
- add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
-
- prompt_embeds = prompt_embeds.to(self.device)
- add_text_embeds = add_text_embeds.to(self.device)
- add_time_ids = add_time_ids.to(self.device).repeat(batch_size * num_images_per_prompt, 1)
-
- # 8. Denoising loop
- for i, t in enumerate(timesteps):
- # Write latents out and skip
- if i < idx_start:
- list_latents_out.append(None)
- continue
- # Set the right starting latents
- elif i == idx_start:
- latents = latents_start.clone()
-
- # Mix latents for crossfeeding
- if i > 0 and list_mixing_coeffs[i] > 0:
- latents_mixtarget = list_latents_mixing[i - 1].clone()
- latents = interpolate_spherical(latents, latents_mixtarget, list_mixing_coeffs[i])
-
-
- # expand the latents if we are doing classifier free guidance
- latent_model_input = torch.cat([latents] * 2)# if do_classifier_free_guidance else latents
- # Always scale latents
- latent_model_input = self.pipe.scheduler.scale_model_input(latent_model_input, t)
-
- # predict the noise residual
- added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
- noise_pred = self.pipe.unet(
- latent_model_input,
- t,
- encoder_hidden_states=prompt_embeds,
- cross_attention_kwargs=cross_attention_kwargs,
- added_cond_kwargs=added_cond_kwargs,
- return_dict=False,
- )[0]
-
- # perform guidance
- if do_classifier_free_guidance:
- noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
- noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
-
- # compute the previous noisy sample x_t -> x_t-1
- latents = self.pipe.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
-
- # Append latents
- list_latents_out.append(latents.clone())
-
- if return_image:
- return self.latent2image(latents)
- else:
- return list_latents_out
#%%
@@ -757,17 +643,7 @@ if __name__ == "__main__":
pipe.vae = AutoencoderTiny.from_pretrained('madebyollin/taesdxl', torch_device='cuda', torch_dtype=torch.float16)
pipe.vae = pipe.vae.cuda()
#%%
- self = DiffusersHolder(pipe)
- self.set_num_inference_steps(4)
- prompt1 = "Photo of a colorful landscape with a blue sky with clouds"
- text_embeddings1 = self.get_text_embedding(prompt1)
- latents_start = self.get_noise(seed=420)
- latents = self.run_diffusion_sd_xl(text_embeddings1, latents_start, idx_start=0, return_image=False)[-1]
- image = self.latent2image(latents)
-
-
-
- xxxx
+
# # xxx
# self.set_dimensions((512, 512))
# self.set_num_inference_steps(4)
@@ -785,6 +661,7 @@ if __name__ == "__main__":
self.set_num_inference_steps(num_inference_steps)
latents_start = self.get_noise()
guidance_scale = 0
+ self.guidance_scale = 0
#% get embeddings1
prompt1 = "Photo of a colorful landscape with a blue sky with clouds"
@@ -797,11 +674,17 @@ if __name__ == "__main__":
prompt_embeds2, negative_prompt_embeds2, pooled_prompt_embeds2, negative_pooled_prompt_embeds2 = text_embeddings2
latents1 = self.run_diffusion_sd_xl(text_embeddings1, latents_start, idx_start=0, return_image=False)
- latents2 = self.run_diffusion_sd_xl(text_embeddings2, latents_start, idx_start=0, return_image=False)
+
+ img1 = self.run_diffusion_sd_xl(text_embeddings1, latents_start, idx_start=0, return_image=True)
+ img1B = self.run_diffusion_sd_xl(text_embeddings1, latents_start, idx_start=0, return_image=True)
- # check if brings same image if restarted
- img1_return = self.run_diffusion_sd_xl(text_embeddings1, latents1[idx_mix-1], idx_start=idx_start, return_image=True)
+
+ # latents2 = self.run_diffusion_sd_xl(text_embeddings2, latents_start, idx_start=0, return_image=False)
+
+
+ # # check if brings same image if restarted
+ # img1_return = self.run_diffusion_sd_xl(text_embeddings1, latents1[idx_mix-1], idx_start=idx_start, return_image=True)
# mix latents
#%%