from copy import deepcopy from dataclasses import dataclass from typing import Any, Callable, Dict, List, Optional, Tuple, Union from diffusers import StableDiffusionXLPipeline from diffusers.image_processor import PipelineImageInput from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img import\ rescale_noise_cfg, retrieve_latents, retrieve_timesteps from diffusers.utils import BaseOutput, deprecate from diffusers.utils.torch_utils import randn_tensor import numpy as np import PIL import torch from ..utils import * from ..utils.sdxl import * BATCH_ORDER = [ "structure_uncond", "appearance_uncond", "uncond", "structure_cond", "appearance_cond", "cond", ] def get_last_control_i(control_schedule, num_inference_steps): if control_schedule is None: return num_inference_steps, num_inference_steps def max_(l): if len(l) == 0: return 0.0 return max(l) structure_max = 0.0 appearance_max = 0.0 for block in control_schedule.values(): if isinstance(block, list): # Handling mid_block block = {0: block} for layer in block.values(): structure_max = max(structure_max, max_(layer[0] + layer[1])) appearance_max = max(appearance_max, max_(layer[2])) structure_i = round(num_inference_steps * structure_max) appearance_i = round(num_inference_steps * appearance_max) return structure_i, appearance_i @dataclass class CtrlXStableDiffusionXLPipelineOutput(BaseOutput): images: Union[List[PIL.Image.Image], np.ndarray] structures = Union[List[PIL.Image.Image], np.ndarray] appearances = Union[List[PIL.Image.Image], np.ndarray] class CtrlXStableDiffusionXLPipeline(StableDiffusionXLPipeline): # diffusers==0.28.0 def prepare_latents( self, image, batch_size, num_images_per_prompt, num_channels_latents, height, width, dtype, device, generator=None, noise=None, ): batch_size = batch_size * num_images_per_prompt if noise is None: shape = ( batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor ) noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) noise = noise * self.scheduler.init_noise_sigma # Starting noise, need to scale else: noise = noise.to(device) if image is None: return noise, None if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): raise ValueError( f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" ) # Offload text encoder if `enable_model_cpu_offload` was enabled if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: self.text_encoder_2.to("cpu") torch.cuda.empty_cache() image = image.to(device=device, dtype=dtype) if image.shape[1] == 4: # Image already in latents form init_latents = image else: # Make sure the VAE is in float32 mode, as it overflows in float16 if self.vae.config.force_upcast: image = image.to(torch.float32) self.vae.to(torch.float32) 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." ) elif isinstance(generator, list): init_latents = [ retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) for i in range(batch_size) ] init_latents = torch.cat(init_latents, dim=0) else: init_latents = retrieve_latents(self.vae.encode(image), generator=generator) if self.vae.config.force_upcast: self.vae.to(dtype) init_latents = init_latents.to(dtype) init_latents = self.vae.config.scaling_factor * init_latents if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: # Expand init_latents for batch_size additional_image_per_prompt = batch_size // init_latents.shape[0] init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: raise ValueError( f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." ) else: init_latents = torch.cat([init_latents], dim=0) return noise, init_latents @property def structure_guidance_scale(self): return self._guidance_scale if self._structure_guidance_scale is None else self._structure_guidance_scale @property def appearance_guidance_scale(self): return self._guidance_scale if self._appearance_guidance_scale is None else self._appearance_guidance_scale @torch.no_grad() def __call__( self, prompt: Union[str, List[str]] = None, # TODO: Support prompt_2 and negative_prompt_2 structure_prompt: Optional[Union[str, List[str]]] = None, appearance_prompt: Optional[Union[str, List[str]]] = None, structure_image: Optional[PipelineImageInput] = None, appearance_image: Optional[PipelineImageInput] = None, num_inference_steps: int = 50, timesteps: List[int] = None, negative_prompt: Optional[Union[str, List[str]]] = None, positive_prompt: Optional[Union[str, List[str]]] = None, height: Optional[int] = None, width: Optional[int] = None, guidance_scale: float = 5.0, structure_guidance_scale: Optional[float] = None, appearance_guidance_scale: Optional[float] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, latents: Optional[torch.Tensor] = None, structure_latents: Optional[torch.Tensor] = None, appearance_latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, # Positive prompt is concatenated with prompt, so no embeddings structure_prompt_embeds: Optional[torch.Tensor] = None, appearance_prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, structure_pooled_prompt_embeds: Optional[torch.Tensor] = None, appearance_pooled_prompt_embeds: Optional[torch.Tensor] = None, negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, control_schedule: Optional[Dict] = None, self_recurrence_schedule: Optional[List[int]] = [], # Format: [(start, end, num_repeat)] decode_structure: Optional[bool] = True, decode_appearance: Optional[bool] = True, output_type: Optional[str] = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, original_size: Tuple[int, int] = None, crops_coords_top_left: Tuple[int, int] = (0, 0), target_size: Tuple[int, int] = None, clip_skip: Optional[int] = None, callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, callback_on_step_end_tensor_inputs: List[str] = ["latents"], **kwargs, ): # TODO: Add function argument documentation callback = kwargs.pop("callback", None) callback_steps = kwargs.pop("callback_steps", None) if callback is not None: deprecate( "callback", "1.0.0", "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", ) if callback_steps is not None: deprecate( "callback_steps", "1.0.0", "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", ) # 0. Default height and width to U-Net height = height or self.default_sample_size * self.vae_scale_factor width = width or self.default_sample_size * self.vae_scale_factor original_size = original_size or (height, width) target_size = target_size or (height, width) # 1. Check inputs. Raise error if not correct self.check_inputs( # TODO: Custom check_inputs for our method prompt, None, # prompt_2 height, width, callback_steps, negative_prompt = negative_prompt, negative_prompt_2 = None, # negative_prompt_2 prompt_embeds = prompt_embeds, negative_prompt_embeds = negative_prompt_embeds, pooled_prompt_embeds = pooled_prompt_embeds, negative_pooled_prompt_embeds = negative_pooled_prompt_embeds, callback_on_step_end_tensor_inputs = callback_on_step_end_tensor_inputs, ) self._guidance_scale = guidance_scale self._structure_guidance_scale = structure_guidance_scale self._appearance_guidance_scale = appearance_guidance_scale self._guidance_rescale = guidance_rescale self._clip_skip = clip_skip self._cross_attention_kwargs = cross_attention_kwargs self._denoising_end = None # denoising_end self._denoising_start = None # denoising_start self._interrupt = False # 2. Define call parameters if prompt is not None and isinstance(prompt, str): batch_size = 1 elif prompt is not None and isinstance(prompt, list): batch_size = len(prompt) else: batch_size = prompt_embeds.shape[0] if batch_size * num_images_per_prompt != 1: raise ValueError( f"Pipeline currently does not support batch_size={batch_size} and num_images_per_prompt=1. " "Effective batch size (batch_size * num_images_per_prompt) must be 1." ) device = self._execution_device # 3. Encode input prompt text_encoder_lora_scale = ( self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None ) if positive_prompt is not None and positive_prompt != "": prompt = prompt + ", " + positive_prompt # Add positive prompt with comma # By default, only add positive prompt to the appearance prompt and not the structure prompt if appearance_prompt is not None and appearance_prompt != "": appearance_prompt = appearance_prompt + ", " + positive_prompt ( prompt_embeds_, negative_prompt_embeds, pooled_prompt_embeds_, negative_pooled_prompt_embeds, ) = self.encode_prompt( prompt = prompt, prompt_2 = None, # prompt_2 device = device, num_images_per_prompt = num_images_per_prompt, do_classifier_free_guidance = True, # self.do_classifier_free_guidance, TODO: Support no CFG negative_prompt = negative_prompt, negative_prompt_2 = None, # negative_prompt_2 prompt_embeds = prompt_embeds, negative_prompt_embeds = negative_prompt_embeds, pooled_prompt_embeds = pooled_prompt_embeds, negative_pooled_prompt_embeds = negative_pooled_prompt_embeds, lora_scale = text_encoder_lora_scale, clip_skip = self.clip_skip, ) prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds_], dim=0).to(device) add_text_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds_], dim=0).to(device) # 3.1. Structure prompt embeddings if structure_prompt is not None and structure_prompt != "": ( structure_prompt_embeds, negative_structure_prompt_embeds, structure_pooled_prompt_embeds, negative_structure_pooled_prompt_embeds, ) = self.encode_prompt( prompt = structure_prompt, prompt_2 = None, # prompt_2 device = device, num_images_per_prompt = num_images_per_prompt, do_classifier_free_guidance = True, # self.do_classifier_free_guidance, TODO: Support no CFG negative_prompt = negative_prompt if structure_image is None else "", negative_prompt_2 = None, # negative_prompt_2 prompt_embeds = structure_prompt_embeds, negative_prompt_embeds = None, # negative_prompt_embeds pooled_prompt_embeds = structure_pooled_prompt_embeds, negative_pooled_prompt_embeds = None, # negative_pooled_prompt_embeds lora_scale = text_encoder_lora_scale, clip_skip = self.clip_skip, ) structure_prompt_embeds = torch.cat( [negative_structure_prompt_embeds, structure_prompt_embeds], dim=0 ).to(device) structure_add_text_embeds = torch.cat( [negative_structure_pooled_prompt_embeds, structure_pooled_prompt_embeds], dim=0 ).to(device) else: structure_prompt_embeds = prompt_embeds structure_add_text_embeds = add_text_embeds # 3.2. Appearance prompt embeddings if appearance_prompt is not None and appearance_prompt != "": ( appearance_prompt_embeds, negative_appearance_prompt_embeds, appearance_pooled_prompt_embeds, negative_appearance_pooled_prompt_embeds, ) = self.encode_prompt( prompt = appearance_prompt, prompt_2 = None, # prompt_2 device = device, num_images_per_prompt = num_images_per_prompt, do_classifier_free_guidance = True, # self.do_classifier_free_guidance, TODO: Support no CFG negative_prompt = negative_prompt if appearance_image is None else "", negative_prompt_2 = None, # negative_prompt_2 prompt_embeds = appearance_prompt_embeds, negative_prompt_embeds = None, # negative_prompt_embeds pooled_prompt_embeds = appearance_pooled_prompt_embeds, # pooled_prompt_embeds negative_pooled_prompt_embeds = None, # negative_pooled_prompt_embeds lora_scale = text_encoder_lora_scale, clip_skip = self.clip_skip, ) appearance_prompt_embeds = torch.cat( [negative_appearance_prompt_embeds, appearance_prompt_embeds], dim=0 ).to(device) appearance_add_text_embeds = torch.cat( [negative_appearance_pooled_prompt_embeds, appearance_pooled_prompt_embeds], dim=0 ).to(device) else: appearance_prompt_embeds = prompt_embeds appearance_add_text_embeds = add_text_embeds # 3.3. Prepare added time ids & embeddings, TODO: Support no CFG if self.text_encoder_2 is None: text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1]) else: text_encoder_projection_dim = self.text_encoder_2.config.projection_dim add_time_ids = self._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 add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0).to(device) # 4. Prepare timesteps timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) # 5. Prepare latent variables num_channels_latents = self.unet.config.in_channels latents, _ = self.prepare_latents( None, batch_size, num_images_per_prompt, num_channels_latents, height, width, prompt_embeds.dtype, device, generator, latents ) if structure_image is not None: structure_image = preprocess( # Center crop + resize structure_image, self.image_processor, height=height, width=width, resize_mode="crop" ) _, clean_structure_latents = self.prepare_latents( structure_image, batch_size, num_images_per_prompt, num_channels_latents, height, width, prompt_embeds.dtype, device, generator, structure_latents, ) else: clean_structure_latents = None structure_latents = latents if structure_latents is None else structure_latents if appearance_image is not None: appearance_image = preprocess( # Center crop + resize appearance_image, self.image_processor, height=height, width=width, resize_mode="crop" ) _, clean_appearance_latents = self.prepare_latents( appearance_image, batch_size, num_images_per_prompt, num_channels_latents, height, width, prompt_embeds.dtype, device, generator, appearance_latents, ) else: clean_appearance_latents = None appearance_latents = latents if appearance_latents is None else appearance_latents # 6. Prepare extra step kwargs extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) # 7. Denoising loop num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) # 7.1 Apply denoising_end def denoising_value_valid(dnv): return isinstance(self.denoising_end, float) and 0 < dnv < 1 if ( self.denoising_end is not None and self.denoising_start is not None and denoising_value_valid(self.denoising_end) and denoising_value_valid(self.denoising_start) and self.denoising_start >= self.denoising_end ): raise ValueError( f"`denoising_start`: {self.denoising_start} cannot be larger than or equal to `denoising_end`: " + f" {self.denoising_end} when using type float." ) elif self.denoising_end is not None and denoising_value_valid(self.denoising_end): discrete_timestep_cutoff = int( round( self.scheduler.config.num_train_timesteps - (self.denoising_end * self.scheduler.config.num_train_timesteps) ) ) num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) timesteps = timesteps[:num_inference_steps] # 7.2 Optionally get guidance scale embedding timestep_cond = None if self.unet.config.time_cond_proj_dim is not None: # TODO: Make guidance scale embedding work with batch_order guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt) timestep_cond = self.get_guidance_scale_embedding( guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim ).to(device=device, dtype=latents.dtype) # 7.3 Get batch order batch_order = deepcopy(BATCH_ORDER) if structure_image is not None: # If image is provided, not generating, so no CFG needed batch_order.remove("structure_uncond") if appearance_image is not None: batch_order.remove("appearance_uncond") structure_control_stop_i, appearance_control_stop_i = get_last_control_i(control_schedule, num_inference_steps) if self_recurrence_schedule is None: self_recurrence_schedule = [0] * num_inference_steps self._num_timesteps = len(timesteps) with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): if self.interrupt: continue if i == structure_control_stop_i: # If not generating structure/appearance, drop after last control if "structure_uncond" not in batch_order: batch_order.remove("structure_cond") if i == appearance_control_stop_i: if "appearance_uncond" not in batch_order: batch_order.remove("appearance_cond") register_attr(self, t=t.item(), do_control=True, batch_order=batch_order) # TODO: For now, assume we are doing classifier-free guidance, support no CF-guidance later latent_model_input = self.scheduler.scale_model_input(latents, t) structure_latent_model_input = self.scheduler.scale_model_input(structure_latents, t) appearance_latent_model_input = self.scheduler.scale_model_input(appearance_latents, t) all_latent_model_input = { "structure_uncond": structure_latent_model_input[0:1], "appearance_uncond": appearance_latent_model_input[0:1], "uncond": latent_model_input[0:1], "structure_cond": structure_latent_model_input[0:1], "appearance_cond": appearance_latent_model_input[0:1], "cond": latent_model_input[0:1], } all_prompt_embeds = { "structure_uncond": structure_prompt_embeds[0:1], "appearance_uncond": appearance_prompt_embeds[0:1], "uncond": prompt_embeds[0:1], "structure_cond": structure_prompt_embeds[1:2], "appearance_cond": appearance_prompt_embeds[1:2], "cond": prompt_embeds[1:2], } all_add_text_embeds = { "structure_uncond": structure_add_text_embeds[0:1], "appearance_uncond": appearance_add_text_embeds[0:1], "uncond": add_text_embeds[0:1], "structure_cond": structure_add_text_embeds[1:2], "appearance_cond": appearance_add_text_embeds[1:2], "cond": add_text_embeds[1:2], } all_time_ids = { "structure_uncond": add_time_ids[0:1], "appearance_uncond": add_time_ids[0:1], "uncond": add_time_ids[0:1], "structure_cond": add_time_ids[1:2], "appearance_cond": add_time_ids[1:2], "cond": add_time_ids[1:2], } concat_latent_model_input = batch_dict_to_tensor(all_latent_model_input, batch_order) concat_prompt_embeds = batch_dict_to_tensor(all_prompt_embeds, batch_order) concat_add_text_embeds = batch_dict_to_tensor(all_add_text_embeds, batch_order) concat_add_time_ids = batch_dict_to_tensor(all_time_ids, batch_order) # Predict the noise residual added_cond_kwargs = {"text_embeds": concat_add_text_embeds, "time_ids": concat_add_time_ids} concat_noise_pred = self.unet( concat_latent_model_input, t, encoder_hidden_states = concat_prompt_embeds, timestep_cond = timestep_cond, cross_attention_kwargs = self.cross_attention_kwargs, added_cond_kwargs = added_cond_kwargs, ).sample all_noise_pred = batch_tensor_to_dict(concat_noise_pred, batch_order) # Classifier-free guidance, TODO: Support no CFG noise_pred = all_noise_pred["uncond"] +\ self.guidance_scale * (all_noise_pred["cond"] - all_noise_pred["uncond"]) structure_noise_pred = all_noise_pred["structure_cond"]\ if "structure_cond" in batch_order else noise_pred if "structure_uncond" in all_noise_pred: structure_noise_pred = all_noise_pred["structure_uncond"] +\ self.structure_guidance_scale * (structure_noise_pred - all_noise_pred["structure_uncond"]) appearance_noise_pred = all_noise_pred["appearance_cond"]\ if "appearance_cond" in batch_order else noise_pred if "appearance_uncond" in all_noise_pred: appearance_noise_pred = all_noise_pred["appearance_uncond"] +\ self.appearance_guidance_scale * (appearance_noise_pred - all_noise_pred["appearance_uncond"]) if self.guidance_rescale > 0.0: noise_pred = rescale_noise_cfg( noise_pred, all_noise_pred["cond"], guidance_rescale=self.guidance_rescale ) if "structure_uncond" in all_noise_pred: structure_noise_pred = rescale_noise_cfg( structure_noise_pred, all_noise_pred["structure_cond"], guidance_rescale=self.guidance_rescale ) if "appearance_uncond" in all_noise_pred: appearance_noise_pred = rescale_noise_cfg( appearance_noise_pred, all_noise_pred["appearance_cond"], guidance_rescale=self.guidance_rescale ) # Compute the previous noisy sample x_t -> x_t-1 concat_noise_pred = torch.cat( [structure_noise_pred, appearance_noise_pred, noise_pred], dim=0, ) concat_latents = torch.cat( [structure_latents, appearance_latents, latents], dim=0, ) structure_latents, appearance_latents, latents = self.scheduler.step( concat_noise_pred, t, concat_latents, **extra_step_kwargs, ).prev_sample.chunk(3) if clean_structure_latents is not None: structure_latents = noise_prev(self.scheduler, t, clean_structure_latents) if clean_appearance_latents is not None: appearance_latents = noise_prev(self.scheduler, t, clean_appearance_latents) # Self-recurrence for _ in range(self_recurrence_schedule[i]): if hasattr(self.scheduler, "_step_index"): # For fancier schedulers self.scheduler._step_index -= 1 # TODO: Does this actually work? t_prev = 0 if i + 1 >= num_inference_steps else timesteps[i + 1] latents = noise_t2t(self.scheduler, t_prev, t, latents) latent_model_input = torch.cat([latents] * 2) register_attr(self, t=t.item(), do_control=False, batch_order=["uncond", "cond"]) # Predict the noise residual added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} noise_pred_uncond, noise_pred_ = self.unet( latent_model_input, t, encoder_hidden_states = prompt_embeds, timestep_cond = timestep_cond, cross_attention_kwargs = self.cross_attention_kwargs, added_cond_kwargs = added_cond_kwargs, ).sample.chunk(2) noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_ - noise_pred_uncond) if self.guidance_rescale > 0.0: noise_pred = rescale_noise_cfg(noise_pred, noise_pred_, guidance_rescale=self.guidance_rescale) latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample # Callbacks if callback_on_step_end is not None: callback_kwargs = {} for k in callback_on_step_end_tensor_inputs: callback_kwargs[k] = locals()[k] callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) latents = callback_outputs.pop("latents", latents) prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) negative_pooled_prompt_embeds = callback_outputs.pop( "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds ) add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) add_neg_time_ids = callback_outputs.pop("add_neg_time_ids", add_neg_time_ids) if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() if callback is not None and i % callback_steps == 0: step_idx = i // getattr(self.scheduler, "order", 1) callback(step_idx, t, latents) # "Reconstruction" if clean_structure_latents is not None: structure_latents = clean_structure_latents if clean_appearance_latents is not None: appearance_latents = clean_appearance_latents # For passing important information onto the refiner self.refiner_args = {"latents": latents.detach(), "prompt": prompt, "negative_prompt": negative_prompt} if not output_type == "latent": # Make sure the VAE is in float32 mode, as it overflows in float16 if self.vae.config.force_upcast: self.vae.to(torch.float32) # self.upcast_vae() is buggy latents = latents.to(torch.float32) structure_latents = structure_latents.to(torch.float32) appearance_latents = appearance_latents.to(torch.float32) image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] image = self.image_processor.postprocess(image, output_type=output_type) if decode_structure: structure = self.vae.decode(structure_latents / self.vae.config.scaling_factor, return_dict=False)[0] structure = self.image_processor.postprocess(structure, output_type=output_type) else: structure = structure_latents if decode_appearance: appearance = self.vae.decode(appearance_latents / self.vae.config.scaling_factor, return_dict=False)[0] appearance = self.image_processor.postprocess(appearance, output_type=output_type) else: appearance = appearance_latents # Cast back to fp16 if needed if self.vae.config.force_upcast: self.vae.to(dtype=torch.float16) else: return CtrlXStableDiffusionXLPipelineOutput( images=latents, structures=structure_latents, appearances=appearance_latents ) # Offload all models self.maybe_free_model_hooks() if not return_dict: return (image, structure, appearance) return CtrlXStableDiffusionXLPipelineOutput(images=image, structures=structure, appearances=appearance)