""" ADOBE CONFIDENTIAL Copyright 2024 Adobe All Rights Reserved. NOTICE: All information contained herein is, and remains the property of Adobe and its suppliers, if any. The intellectual and technical concepts contained herein are proprietary to Adobe and its suppliers and are protected by all applicable intellectual property laws, including trade secret and copyright laws. Dissemination of this information or reproduction of this material is strictly forbidden unless prior written permission is obtained from Adobe. """ from typing import Callable, List, Optional, Union import inspect import einops import PIL.Image import numpy as np import torch as th import torch.nn as nn from torchvision import transforms from diffusers import ModelMixin from transformers import AutoModel, AutoConfig, SiglipVisionConfig, Dinov2Config, Dinov2Model from transformers import SiglipVisionModel from diffusers import DiffusionPipeline from diffusers.image_processor import VaeImageProcessor from diffusers.models import AutoencoderKL, UNet2DConditionModel from diffusers.schedulers import KarrasDiffusionSchedulers from diffusers.utils.torch_utils import randn_tensor from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput from diffusers.configuration_utils import ConfigMixin, register_to_config # REf: https://github.com/tatp22/multidim-positional-encoding/tree/master OUT_SIZE = 768 IN_SIZE = 2048 DINO_SIZE = 224 DINO_MEAN = [0.485, 0.456, 0.406] DINO_STD = [0.229, 0.224, 0.225] SIGLIP_SIZE = 256 SIGLIP_MEAN = [0.5] SIGLIP_STD = [0.5] def get_emb(sin_inp): """ Gets a base embedding for one dimension with sin and cos intertwined """ emb = th.stack((sin_inp.sin(), sin_inp.cos()), dim=-1) return th.flatten(emb, -2, -1) class PositionalEncoding1D(nn.Module): def __init__(self, channels): """ :param channels: The last dimension of the tensor you want to apply pos emb to. """ super(PositionalEncoding1D, self).__init__() self.org_channels = channels channels = int(np.ceil(channels / 2) * 2) self.channels = channels inv_freq = 1.0 / (10000 ** (th.arange(0, channels, 2).float() / channels)) self.register_buffer("inv_freq", inv_freq) self.register_buffer("cached_penc", None, persistent=False) def forward(self, tensor): """ :param tensor: A 3d tensor of size (batch_size, x, ch) :return: Positional Encoding Matrix of size (batch_size, x, ch) """ if len(tensor.shape) != 3: raise RuntimeError("The input tensor has to be 3d!") if self.cached_penc is not None and self.cached_penc.shape == tensor.shape: return self.cached_penc self.cached_penc = None batch_size, x, orig_ch = tensor.shape pos_x = th.arange(x, device=tensor.device, dtype=self.inv_freq.dtype) sin_inp_x = th.einsum("i,j->ij", pos_x, self.inv_freq) emb_x = get_emb(sin_inp_x) emb = th.zeros((x, self.channels), device=tensor.device, dtype=tensor.dtype) emb[:, : self.channels] = emb_x self.cached_penc = emb[None, :, :orig_ch].repeat(batch_size, 1, 1) return self.cached_penc class PositionalEncoding3D(nn.Module): def __init__(self, channels): """ :param channels: The last dimension of the tensor you want to apply pos emb to. """ super(PositionalEncoding3D, self).__init__() self.org_channels = channels channels = int(np.ceil(channels / 6) * 2) if channels % 2: channels += 1 self.channels = channels inv_freq = 1.0 / (10000 ** (th.arange(0, channels, 2).float() / channels)) self.register_buffer("inv_freq", inv_freq) self.register_buffer("cached_penc", None, persistent=False) def forward(self, tensor): """ :param tensor: A 5d tensor of size (batch_size, x, y, z, ch) :return: Positional Encoding Matrix of size (batch_size, x, y, z, ch) """ if len(tensor.shape) != 5: raise RuntimeError("The input tensor has to be 5d!") if self.cached_penc is not None and self.cached_penc.shape == tensor.shape: return self.cached_penc self.cached_penc = None batch_size, x, y, z, orig_ch = tensor.shape pos_x = th.arange(x, device=tensor.device, dtype=self.inv_freq.dtype) pos_y = th.arange(y, device=tensor.device, dtype=self.inv_freq.dtype) pos_z = th.arange(z, device=tensor.device, dtype=self.inv_freq.dtype) sin_inp_x = th.einsum("i,j->ij", pos_x, self.inv_freq) sin_inp_y = th.einsum("i,j->ij", pos_y, self.inv_freq) sin_inp_z = th.einsum("i,j->ij", pos_z, self.inv_freq) emb_x = get_emb(sin_inp_x).unsqueeze(1).unsqueeze(1) emb_y = get_emb(sin_inp_y).unsqueeze(1) emb_z = get_emb(sin_inp_z) emb = th.zeros( (x, y, z, self.channels * 3), device=tensor.device, dtype=tensor.dtype, ) emb[:, :, :, : self.channels] = emb_x emb[:, :, :, self.channels : 2 * self.channels] = emb_y emb[:, :, :, 2 * self.channels :] = emb_z self.cached_penc = emb[None, :, :, :, :orig_ch].repeat(batch_size, 1, 1, 1, 1) return self.cached_penc class AnalogyInputProcessor(ModelMixin, ConfigMixin): @register_to_config def __init__(self,): super(AnalogyInputProcessor, self).__init__() self.dino_transform = transforms.Compose( [ transforms.Resize((DINO_SIZE, DINO_SIZE)), transforms.ToTensor(), transforms.Normalize(DINO_MEAN, DINO_STD), # SIGLIP normalization ] ) self.siglip_transform = transforms.Compose( [ transforms.Resize((SIGLIP_SIZE, SIGLIP_SIZE)), transforms.ToTensor(), transforms.Normalize(SIGLIP_MEAN, SIGLIP_STD), # SIGLIP normalization ] ) dino_mean = th.tensor(DINO_MEAN).view(1, 3, 1, 1) dino_std = th.tensor(DINO_STD).view(1, 3, 1, 1) siglip_mean = [SIGLIP_MEAN[0],] * 3 siglip_std = [SIGLIP_STD[0],] * 3 siglip_mean = th.tensor(siglip_mean).view(1, 3, 1, 1) siglip_std = th.tensor(siglip_std).view(1, 3, 1, 1) self.register_buffer("dino_mean", dino_mean) self.register_buffer("dino_std", dino_std) self.register_buffer("siglip_mean", siglip_mean) self.register_buffer("siglip_std", siglip_std) def __call__(self, analogy_prompt): # List of tuples of (A, A*, B) img_a_dino = [] img_a_siglip = [] img_a_star_dino = [] img_a_star_siglip = [] img_b_dino = [] img_b_siglip = [] for im_set in analogy_prompt: img_a, img_a_star, img_b = im_set img_a_dino.append(self.dino_transform(img_a)) img_a_siglip.append(self.siglip_transform(img_a)) img_a_star_dino.append(self.dino_transform(img_a_star)) img_a_star_siglip.append(self.siglip_transform(img_a_star)) img_b_dino.append(self.dino_transform(img_b)) img_b_siglip.append(self.siglip_transform(img_b)) img_a_dino = th.stack(img_a_dino, 0) img_a_siglip = th.stack(img_a_siglip, 0) img_a_star_dino = th.stack(img_a_star_dino, 0) img_a_star_siglip = th.stack(img_a_star_siglip, 0) img_b_dino = th.stack(img_b_dino, 0) img_b_siglip = th.stack(img_b_siglip, 0) dino_combined_input = th.stack([img_b_dino, img_a_dino, img_a_star_dino], 0) siglip_combined_input = th.stack([img_b_siglip, img_a_siglip, img_a_star_siglip], 0) return dino_combined_input, siglip_combined_input def get_negative(self, dino_in, siglip_in): dino_i = ((dino_in * 0 + 0.5) - self.dino_mean) / self.dino_std siglip_i = ((siglip_in * 0 + 0.5) - self.siglip_mean) / self.siglip_std return dino_i, siglip_i class AnalogyProjector(ModelMixin, ConfigMixin): @register_to_config def __init__(self): super(AnalogyProjector, self).__init__() self.projector = DinoSiglipMixer() self.pos_embd_1D = PositionalEncoding1D(OUT_SIZE) self.pos_embd_3D = PositionalEncoding3D(OUT_SIZE) def forward(self, dino_in, siglip_in, batch_size): image_embeddings = self.projector(dino_in, siglip_in) image_embeddings = einops.rearrange(image_embeddings, '(k b) t d -> b k t d', b=batch_size) image_embeddings = self.position_embd(image_embeddings) return image_embeddings def position_embd(self, image_embeddings, concat=False): canvas_embd = image_embeddings[:, :, 1:, :] batch_size = canvas_embd.shape[0] type_size = canvas_embd.shape[1] xy_size = canvas_embd.shape[2] x_size = int(xy_size ** 0.5) canvas_embd = canvas_embd.reshape(batch_size, type_size, x_size, x_size, -1) if concat: canvas_embd = th.cat([canvas_embd, self.pos_embd_3D(canvas_embd)], -1) else: canvas_embd = self.pos_embd_3D(canvas_embd) + canvas_embd canvas_embd = canvas_embd.reshape(batch_size, type_size, xy_size, -1) class_embd = image_embeddings[:, :, 0, :] if concat: class_embd = th.cat([class_embd, self.pos_embd_1D(class_embd)], -1) else: class_embd = self.pos_embd_1D(class_embd) + class_embd all_embd_list = [] for i in range(type_size): all_embd_list.append(class_embd[:, i:i+1]) all_embd_list.append(canvas_embd[:, i]) image_embeddings = th.cat(all_embd_list, 1) return image_embeddings class HighLowMixer(th.nn.Module): def __init__(self, in_size=IN_SIZE, out_size=OUT_SIZE): super().__init__() mid_size = (in_size + out_size) // 2 self.lower_projector = th.nn.Sequential( th.nn.LayerNorm(IN_SIZE//2), th.nn.SiLU() ) self.upper_projector = th.nn.Sequential( th.nn.LayerNorm(IN_SIZE//2), th.nn.SiLU() ) self.projectors = th.nn.ModuleList([ # add layer norm th.nn.Linear(in_size, mid_size), th.nn.SiLU(), th.nn.Linear(mid_size, out_size) ]) # initialize for proj in self.projectors: if isinstance(proj, th.nn.Linear): th.nn.init.xavier_uniform_(proj.weight) th.nn.init.zeros_(proj.bias) def forward(self, lower_in, upper_in, ): # ALso format lower_in lower_in = self.lower_projector(lower_in) upper_in = self.upper_projector(upper_in) x = th.cat([lower_in, upper_in], -1) for proj in self.projectors: x = proj(x) return x class DinoSiglipMixer(th.nn.Module): def __init__(self, in_size=OUT_SIZE * 2, out_size=OUT_SIZE): super().__init__() self.dino_projector = HighLowMixer() self.siglip_projector = HighLowMixer() self.projectors = th.nn.Sequential( th.nn.SiLU(), th.nn.Linear(in_size, out_size), ) # initialize for proj in self.projectors: if isinstance(proj, th.nn.Linear): th.nn.init.xavier_uniform_(proj.weight) th.nn.init.zeros_(proj.bias) def forward(self, dino_in, siglip_in): # ALso format lower_in lower, upper = th.chunk(dino_in, 2, -1) dino_out = self.dino_projector(lower, upper) lower, upper = th.chunk(siglip_in, 2, -1) siglip_out = self.siglip_projector(lower, upper) x = th.cat([dino_out, siglip_out], -1) for proj in self.projectors: x = proj(x) return x class AnalogyEncoder(ModelMixin, ConfigMixin): @register_to_config def __init__(self, load_pretrained=False, dino_config_dict=None, siglip_config_dict=None): super().__init__() if load_pretrained: image_encoder_dino = AutoModel.from_pretrained('facebook/dinov2-large', torch_dtype=th.float16) image_encoder_siglip = SiglipVisionModel.from_pretrained("google/siglip-large-patch16-256", torch_dtype=th.float16, attn_implementation="sdpa") else: image_encoder_dino = AutoModel.from_config(Dinov2Config.from_dict(dino_config_dict)) image_encoder_siglip = AutoModel.from_config(SiglipVisionConfig.from_dict(siglip_config_dict)) image_encoder_dino.requires_grad_(False) image_encoder_dino = image_encoder_dino.to(memory_format=th.channels_last) image_encoder_siglip.requires_grad_(False) image_encoder_siglip = image_encoder_siglip.to(memory_format=th.channels_last) self.image_encoder_dino = image_encoder_dino self.image_encoder_siglip = image_encoder_siglip def dino_normalization(self, encoder_output): embeds = encoder_output.last_hidden_state embeds_pooled = embeds[:, 0:1] embeds = embeds / th.norm(embeds_pooled, dim=-1, keepdim=True) return embeds def siglip_normalization(self, encoder_output): embeds = th.cat ([encoder_output.pooler_output[:, None, :], encoder_output.last_hidden_state], dim=1) embeds_pooled = embeds[:, 0:1] embeds = embeds / th.norm(embeds_pooled, dim=-1, keepdim=True) return embeds def forward(self, dino_in, siglip_in): x_1 = self.image_encoder_dino(dino_in, output_hidden_states=True) x_1_first = x_1.hidden_states[0] x_1 = self.dino_normalization(x_1) x_2 = self.image_encoder_siglip(siglip_in, output_hidden_states=True) x_2_first = x_2.hidden_states[0] x_2_first_pool = th.mean(x_2_first, dim=1, keepdim=True) x_2_first = th.cat([x_2_first_pool, x_2_first], 1) x_2 = self.siglip_normalization(x_2) dino_embd = th.cat([x_1, x_1_first], -1) siglip_embd = th.cat([x_2, x_2_first], -1) return dino_embd, siglip_embd class PatternAnalogyTrifuser(DiffusionPipeline): r""" This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) """ model_cpu_offload_seq = "bert->unet->vqvae" analogy_input_processor: AnalogyInputProcessor analogy_encoder: AnalogyEncoder analogy_projector: AnalogyProjector unet: UNet2DConditionModel vae: AutoencoderKL scheduler: KarrasDiffusionSchedulers def __init__(self, analogy_input_processor: AnalogyInputProcessor, analogy_projector: AnalogyProjector, analogy_encoder: AnalogyEncoder, unet: UNet2DConditionModel, vae: AutoencoderKL, scheduler: KarrasDiffusionSchedulers,): super().__init__() self.register_modules( analogy_input_processor=analogy_input_processor, analogy_encoder=analogy_encoder, analogy_projector=analogy_projector, unet=unet, vae=vae, scheduler=scheduler, ) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_image_variation.StableDiffusionImageVariationPipeline.check_inputs def check_inputs(self, analogy_prompt, negative_analogy_prompt, height, width, callback_steps): if ( not isinstance(analogy_prompt, th.Tensor) and not isinstance(analogy_prompt, PIL.Image.Image) and not isinstance(analogy_prompt, list) ): raise ValueError( "`analogy_prompt` contents have to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is" f" {type(analogy_prompt)}" ) if not negative_analogy_prompt is None: if ( not isinstance(negative_analogy_prompt, th.Tensor) and not isinstance(negative_analogy_prompt, PIL.Image.Image) and not isinstance(negative_analogy_prompt, list) ): raise ValueError( "`negative_analogy_prompt` contents have to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is" f" {type(negative_analogy_prompt)}" ) if height % 8 != 0 or width % 8 != 0: raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") if (callback_steps is None) or ( callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) ): raise ValueError( f"`callback_steps` has to be a positive integer but is {callback_steps} of type" f" {type(callback_steps)}." ) # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None): shape = ( batch_size, num_channels_latents, int(height) // self.vae_scale_factor, int(width) // self.vae_scale_factor, ) if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" f" size of {batch_size}. Make sure the batch size matches the length of the generators." ) if latents is None: latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) else: latents = latents.to(device) # scale the initial noise by the standard deviation required by the scheduler latents = latents * self.scheduler.init_noise_sigma return latents # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs def prepare_extra_step_kwargs(self, generator, eta): # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 # and should be between [0, 1] accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) extra_step_kwargs = {} if accepts_eta: extra_step_kwargs["eta"] = eta # check if the scheduler accepts generator accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) if accepts_generator: extra_step_kwargs["generator"] = generator return extra_step_kwargs # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None): shape = ( batch_size, num_channels_latents, int(height) // self.vae_scale_factor, int(width) // self.vae_scale_factor, ) if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" f" size of {batch_size}. Make sure the batch size matches the length of the generators." ) if latents is None: latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) else: latents = latents.to(device) # scale the initial noise by the standard deviation required by the scheduler latents = latents * self.scheduler.init_noise_sigma return latents def _encode_prompt(self, analogy_prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt): r""" Encodes the prompt into text encoder hidden states. Args: prompt (`str` or `List[str]`): prompt to be encoded device: (`torch.device`): torch device num_images_per_prompt (`int`): number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not negative_prompt (`str` or `List[str]`): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). """ weight_dtype = self.unet.dtype dino_input, siglip_input = self.analogy_input_processor(analogy_prompt) dino_input = dino_input.to(device=device).to(dtype=weight_dtype) siglip_input = siglip_input.to(device=device).to(dtype=weight_dtype) batch_size = dino_input.shape[1] dino_input_reshaped = einops.rearrange(dino_input, "k b c h w -> (k b) c h w") siglip_input_reshaped = einops.rearrange(siglip_input, "k b c h w -> (k b) c h w") dino_enc, siglip_enc = self.analogy_encoder(dino_input_reshaped, siglip_input_reshaped) image_embeddings = self.analogy_projector(dino_enc, siglip_enc, batch_size) # Check size here. bs_embed, seq_len, _ = image_embeddings.shape image_embeddings = image_embeddings.repeat(num_images_per_prompt, 1, 1) # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance: uncond_images: List[str] if negative_prompt is None: uncond_images = [np.zeros((512, 512, 3)) + 0.5] * batch_size elif type(negative_prompt) is not type(analogy_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(analogy_prompt)} !=" f" {type(negative_prompt)}." ) elif isinstance(negative_prompt, PIL.Image.Image): uncond_images = [negative_prompt] elif batch_size != len(negative_prompt): raise ValueError( f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" f" {analogy_prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" " the batch size of `prompt`." ) else: uncond_images = negative_prompt dino_neg, siglip_neg = self.analogy_input_processor.get_negative(dino_input, siglip_input) dino_neg = dino_neg.to(device=device).to(dtype=weight_dtype) siglip_neg = siglip_neg.to(device=device).to(dtype=weight_dtype) dino_neg_reshaped = einops.rearrange(dino_neg, "k b c h w -> (k b) c h w") siglip_neg_reshaped = einops.rearrange(siglip_neg, "k b c h w -> (k b) c h w") dino_neg_enc, siglip_neg_enc = self.analogy_encoder(dino_neg_reshaped, siglip_neg_reshaped) negative_prompt_embeds = self.analogy_projector(dino_neg_enc, siglip_neg_enc, batch_size) negative_prompt_embeds = negative_prompt_embeds.repeat(num_images_per_prompt, 1, 1) image_embeddings = th.cat([negative_prompt_embeds, image_embeddings]) return image_embeddings @th.no_grad() def __call__( self, analogy_prompt: Union[str, List[str]] = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, height: Optional[int] = None, width: Optional[int] = None, negative_analogy_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, generator: Optional[Union[th.Generator, List[th.Generator]]] = None, latents: Optional[th.FloatTensor] = None, output_type: Optional[str] = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, th.Tensor], None]] = None, callback_steps: int = 1, start_step: int = 0, ): r""" The call function to the pipeline for generation. Args: image (`PIL.Image.Image`, `List[PIL.Image.Image]` or `torch.Tensor`): The image prompt or prompts to guide the image generation. height (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. width (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`): The width in pixels of the generated image. num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. negative_prompt (`str` or `List[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. generator (`torch.Generator`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image generation. Can be used to tweak the same generation with different prompts. If not provided, a latents tensor is generated by sampling using the supplied random `generator`. output_type (`str`, *optional*, defaults to `"pil"`): The output format of the generated image. Choose between `PIL.Image` or `np.array`. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a plain tuple. callback (`Callable`, *optional*): A function that calls every `callback_steps` steps during inference. The function is called with the following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`. callback_steps (`int`, *optional*, defaults to 1): The frequency at which the `callback` function is called. If not specified, the callback is called at every step. Examples: ```py >>> from diffusers import VersatileDiffusionImageVariationPipeline >>> import torch >>> import requests >>> from io import BytesIO >>> from PIL import Image >>> # let's download an initial image >>> url = "https://huggingface.co./datasets/diffusers/images/resolve/main/benz.jpg" >>> response = requests.get(url) >>> image = Image.open(BytesIO(response.content)).convert("RGB") >>> pipe = VersatileDiffusionImageVariationPipeline.from_pretrained( ... "shi-labs/versatile-diffusion", torch_dtype=torch.float16 ... ) >>> pipe = pipe.to("cuda") >>> generator = torch.Generator(device="cuda").manual_seed(0) >>> image = pipe(image, generator=generator).images[0] >>> image.save("./car_variation.png") ``` Returns: [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned, otherwise a `tuple` is returned where the first element is a list with the generated images. """ # 1. Check inputs. Raise error if not correct height = height or self.unet.config.sample_size * self.vae_scale_factor width = width or self.unet.config.sample_size * self.vae_scale_factor # 1. Check inputs. Raise error if not correct self.check_inputs(analogy_prompt, negative_analogy_prompt, height, width, callback_steps) # 2. Define call parameters if isinstance(analogy_prompt, list): batch_size = len(analogy_prompt) elif isinstance(analogy_prompt, tuple): batch_size = 1 else: raise ValueError( f"`analogy_prompt` has to be a list of images or a tuple of images but is of type {type(analogy_prompt)}" ) device = self._execution_device # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` # corresponds to doing no classifier free guidance. do_classifier_free_guidance = guidance_scale > 1.0 # 3. Encode input prompt analogy_embeddings = self._encode_prompt( analogy_prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_analogy_prompt ) # 4. Prepare timesteps self.scheduler.set_timesteps(num_inference_steps, device=device) timesteps = self.scheduler.timesteps # Now this should be from start step onwards timesteps = timesteps[start_step:] # 5. Prepare latent variables num_channels_latents = self.unet.config.in_channels latents = self.prepare_latents( batch_size * num_images_per_prompt, num_channels_latents, height, width, analogy_embeddings.dtype, device, generator, latents, ) # 6. Prepare extra step kwargs. extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) # 7. Denoising loop for i, t in enumerate(self.progress_bar(timesteps)): # expand the latents if we are doing classifier free guidance latent_model_input = th.cat([latents] * 2) if do_classifier_free_guidance else latents latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) # predict the noise residual noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=analogy_embeddings).sample # perform guidance if do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample # call the callback, if provided if callback is not None and i % callback_steps == 0: step_idx = i // getattr(self.scheduler, "order", 1) callback(step_idx, t, latents) if not output_type == "latent": image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] else: image = latents image = self.image_processor.postprocess(image, output_type=output_type) if not return_dict: return (image,) return ImagePipelineOutput(images=image)