from typing import Any, Callable, Dict, List, Optional, Union, Tuple import cv2 import PIL import numpy as np from PIL import Image import torch from torchvision import transforms from insightface.app import FaceAnalysis ### insight-face installation can be found at https://github.com/deepinsight/insightface from safetensors import safe_open from huggingface_hub.utils import validate_hf_hub_args from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput from diffusers.pipelines.stable_diffusion import StableDiffusionPipeline from diffusers.utils.import_utils import is_xformers_available from functions import insert_markers_to_prompt, masks_for_unique_values, apply_mask_to_raw_image, tokenize_and_mask_noun_phrases_ends, prepare_image_token_idx from functions import ProjPlusModel, masks_for_unique_values from attention import Consistent_IPAttProcessor, Consistent_AttProcessor, FacialEncoder from easydict import EasyDict as edict from huggingface_hub import hf_hub_download ### Model can be imported from https://github.com/zllrunning/face-parsing.PyTorch?tab=readme-ov-file ### We use the ckpt of 79999_iter.pth: https://drive.google.com/open?id=154JgKpzCPW82qINcVieuPH3fZ2e0P812 ### Thanks for the open source of face-parsing model. from BiSeNet.model import BiSeNet import os PipelineImageInput = Union[ PIL.Image.Image, torch.FloatTensor, List[PIL.Image.Image], List[torch.FloatTensor], ] ### Download the pretrained model from huggingface and put it locally, then place the model in a local directory and specify the directory location. class ConsistentIDPipeline(StableDiffusionPipeline): def cuda(self, dtype=torch.float16, use_xformers=False): self.to('cuda', dtype) # if hasattr(self, 'image_proj_model'): # self.image_proj_model.to(self.unet.device).to(self.unet.dtype) if use_xformers: if is_xformers_available(): import xformers from packaging import version xformers_version = version.parse(xformers.__version__) if xformers_version == version.parse("0.0.16"): logger.warn( "xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details." ) self.enable_xformers_memory_efficient_attention() else: raise ValueError("xformers is not available. Make sure it is installed correctly") @validate_hf_hub_args def load_ConsistentID_model( self, consistentID_weight_path: str, bise_net_weight_path: str, trigger_word_facial: str = '<|facial|>', # A CLIP ViT-H/14 model trained with the LAION-2B English subset of LAION-5B using OpenCLIP. # output dim: 1280. image_encoder_path: str = 'laion/CLIP-ViT-H-14-laion2B-s32B-b79K', torch_dtype = torch.float16, num_tokens = 4, lora_rank= 128, **kwargs, ): self.lora_rank = lora_rank self.torch_dtype = torch_dtype self.num_tokens = num_tokens self.set_ip_adapter() self.image_encoder_path = image_encoder_path self.clip_encoder = CLIPVisionModelWithProjection.from_pretrained(self.image_encoder_path).to( self.device, dtype=self.torch_dtype ) self.clip_preprocessor = CLIPImageProcessor() self.id_image_processor = CLIPImageProcessor() self.crop_size = 512 # FaceID self.app = FaceAnalysis(name="buffalo_l", providers=['CPUExecutionProvider']) self.app.prepare(ctx_id=0, det_size=(640, 640)) if not os.path.exists(consistentID_weight_path): ### Download pretrained models hf_hub_download(repo_id="JackAILab/ConsistentID", repo_type="model", filename=os.path.basename(consistentID_weight_path), local_dir=os.path.dirname(consistentID_weight_path)) if not os.path.exists(bise_net_weight_path): hf_hub_download(repo_id="JackAILab/ConsistentID", filename=os.path.basename(bise_net_weight_path), local_dir=os.path.dirname(bise_net_weight_path)) bise_net = BiSeNet(n_classes = 19) bise_net.load_state_dict(torch.load(bise_net_weight_path, map_location="cpu")) bise_net.to(self.device, dtype=self.torch_dtype) bise_net.eval() self.bise_net = bise_net # Colors for all 20 parts self.part_colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 0, 85], [255, 0, 170], [0, 255, 0], [85, 255, 0], [170, 255, 0], [0, 255, 85], [0, 255, 170], [0, 0, 255], [85, 0, 255], [170, 0, 255], [0, 85, 255], [0, 170, 255], [255, 255, 0], [255, 255, 85], [255, 255, 170], [255, 0, 255], [255, 85, 255], [255, 170, 255], [0, 255, 255], [85, 255, 255], [170, 255, 255]] # image_proj_model maps 1280-dim OpenCLIP embeddings to 768-dim face prompt embeddings. self.image_proj_model = ProjPlusModel( cross_attention_dim=self.unet.config.cross_attention_dim, id_embeddings_dim=512, clip_embeddings_dim=self.clip_encoder.config.hidden_size, num_tokens=self.num_tokens, # 4 - inspirsed by IPAdapter and Midjourney ).to(self.device, dtype=self.torch_dtype) self.FacialEncoder = FacialEncoder().to(self.device, dtype=self.torch_dtype) if consistentID_weight_path.endswith(".safetensors"): state_dict = {"id_encoder": {}, "lora_weights": {}} with safe_open(consistentID_weight_path, framework="pt", device="cpu") as f: ### TODO safetensors add for key in f.keys(): if key.startswith("FacialEncoder."): state_dict["FacialEncoder"][key.replace("FacialEncoder.", "")] = f.get_tensor(key) elif key.startswith("image_proj."): state_dict["image_proj"][key.replace("image_proj.", "")] = f.get_tensor(key) else: state_dict = torch.load(consistentID_weight_path, map_location="cpu") self.trigger_word_facial = trigger_word_facial self.FacialEncoder.load_state_dict(state_dict["FacialEncoder"], strict=True) self.image_proj_model.load_state_dict(state_dict["image_proj"], strict=True) ip_layers = torch.nn.ModuleList(self.unet.attn_processors.values()) ip_layers.load_state_dict(state_dict["adapter_modules"], strict=True) print(f"Successfully loaded weights from checkpoint") # Add trigger word token if self.tokenizer is not None: self.tokenizer.add_tokens([self.trigger_word_facial], special_tokens=True) def set_ip_adapter(self): unet = self.unet attn_procs = {} for name in unet.attn_processors.keys(): cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim if name.startswith("mid_block"): hidden_size = unet.config.block_out_channels[-1] elif name.startswith("up_blocks"): block_id = int(name[len("up_blocks.")]) hidden_size = list(reversed(unet.config.block_out_channels))[block_id] elif name.startswith("down_blocks"): block_id = int(name[len("down_blocks.")]) hidden_size = unet.config.block_out_channels[block_id] if cross_attention_dim is None: attn_procs[name] = Consistent_AttProcessor( hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, rank=self.lora_rank, ).to(self.device, dtype=self.torch_dtype) else: attn_procs[name] = Consistent_IPAttProcessor( hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, scale=1.0, rank=self.lora_rank, num_tokens=self.num_tokens, ).to(self.device, dtype=self.torch_dtype) unet.set_attn_processor(attn_procs) @torch.inference_mode() # parsed_image_parts2 is a batched tensor of parsed_image_parts with bs=1. It only contains the facial areas of one input image. # clip_encoder maps image parts to image-space diffusion prompts. # Then the facial class token embeddings are replaced with the fused (multi_facial_embeds, prompt_embeds[class_tokens_mask]). def get_local_facial_embeds(self, prompt_embeds, uncond_prompt_embeds, parsed_image_parts2, facial_token_masks, valid_facial_token_idx_mask, calc_uncond=True): hidden_states = [] uncond_hidden_states = [] for parsed_image_parts in parsed_image_parts2: hidden_state = self.clip_encoder(parsed_image_parts.to(self.device, dtype=self.torch_dtype), output_hidden_states=True).hidden_states[-2] uncond_hidden_state = self.clip_encoder(torch.zeros_like(parsed_image_parts, dtype=self.torch_dtype).to(self.device), output_hidden_states=True).hidden_states[-2] hidden_states.append(hidden_state) uncond_hidden_states.append(uncond_hidden_state) multi_facial_embeds = torch.stack(hidden_states) uncond_multi_facial_embeds = torch.stack(uncond_hidden_states) # conditional prompt. # FacialEncoder maps multi_facial_embeds to facial ID embeddings, and replaces the class tokens in prompt_embeds # with the fused (facial ID embeddings, prompt_embeds[class_tokens_mask]). # multi_facial_embeds: [1, 5, 257, 1280]. facial_prompt_embeds = self.FacialEncoder(prompt_embeds, multi_facial_embeds, facial_token_masks, valid_facial_token_idx_mask) if not calc_uncond: return facial_prompt_embeds, None # unconditional prompt. uncond_facial_prompt_embeds = self.FacialEncoder(uncond_prompt_embeds, uncond_multi_facial_embeds, facial_token_masks, valid_facial_token_idx_mask) return facial_prompt_embeds, uncond_facial_prompt_embeds @torch.inference_mode() # Extrat OpenCLIP embeddings from the input image and map them to face prompt embeddings. def get_global_id_embeds(self, faceid_embeds, face_image, s_scale, shortcut=False): clip_image = self.clip_preprocessor(images=face_image, return_tensors="pt").pixel_values clip_image = clip_image.to(self.device, dtype=self.torch_dtype) clip_image_embeds = self.clip_encoder(clip_image, output_hidden_states=True).hidden_states[-2] uncond_clip_image_embeds = self.clip_encoder(torch.zeros_like(clip_image), output_hidden_states=True).hidden_states[-2] faceid_embeds = faceid_embeds.to(self.device, dtype=self.torch_dtype) # image_proj_model maps 1280-dim OpenCLIP embeddings to 768-dim face prompt embeddings. # clip_image_embeds are used as queries to transform faceid_embeds. # faceid_embeds -> kv, clip_image_embeds -> q global_id_embeds = self.image_proj_model(faceid_embeds, clip_image_embeds, shortcut=shortcut, scale=s_scale) uncond_global_id_embeds = self.image_proj_model(torch.zeros_like(faceid_embeds), uncond_clip_image_embeds, shortcut=shortcut, scale=s_scale) return global_id_embeds, uncond_global_id_embeds def set_scale(self, scale): for attn_processor in self.pipe.unet.attn_processors.values(): if isinstance(attn_processor, Consistent_IPAttProcessor): attn_processor.scale = scale @torch.inference_mode() def extract_faceid(self, face_image): faceid_image = np.array(face_image) faces = self.app.get(faceid_image) if faces==[]: faceid_embeds = torch.zeros_like(torch.empty((1, 512))) else: faceid_embeds = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0) return faceid_embeds @torch.inference_mode() def parse_face_mask(self, raw_image_refer): to_tensor = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), ]) to_pil = transforms.ToPILImage() with torch.no_grad(): image = raw_image_refer.resize((512, 512), Image.BILINEAR) image_resize_PIL = image img = to_tensor(image) img = torch.unsqueeze(img, 0) img = img.to(self.device, dtype=self.torch_dtype) out = self.bise_net(img)[0] parsing_anno = out.squeeze(0).cpu().numpy().argmax(0) im = np.array(image_resize_PIL) vis_im = im.copy().astype(np.uint8) stride=1 vis_parsing_anno = parsing_anno.copy().astype(np.uint8) vis_parsing_anno = cv2.resize(vis_parsing_anno, None, fx=stride, fy=stride, interpolation=cv2.INTER_NEAREST) vis_parsing_anno_color = np.zeros((vis_parsing_anno.shape[0], vis_parsing_anno.shape[1], 3)) + 255 num_of_class = np.max(vis_parsing_anno) for pi in range(1, num_of_class + 1): # num_of_class=17 pi=1~16 index = np.where(vis_parsing_anno == pi) vis_parsing_anno_color[index[0], index[1], :] = self.part_colors[pi] vis_parsing_anno_color = vis_parsing_anno_color.astype(np.uint8) vis_parsing_anno_color = cv2.addWeighted(cv2.cvtColor(vis_im, cv2.COLOR_RGB2BGR), 0.4, vis_parsing_anno_color, 0.6, 0) return vis_parsing_anno_color, vis_parsing_anno @torch.inference_mode() def extract_facemask(self, input_image_obj): vis_parsing_anno_color, vis_parsing_anno = self.parse_face_mask(input_image_obj) parsing_mask_list = masks_for_unique_values(vis_parsing_anno) key_parsing_mask_dict = {} key_list = ["Face", "Left_Ear", "Right_Ear", "Left_Eye", "Right_Eye", "Nose", "Upper_Lip", "Lower_Lip"] processed_keys = set() for key, mask_image in parsing_mask_list.items(): if key in key_list: if "_" in key: prefix = key.split("_")[1] if prefix in processed_keys: continue else: key_parsing_mask_dict[key] = mask_image processed_keys.add(prefix) key_parsing_mask_dict[key] = mask_image return key_parsing_mask_dict, vis_parsing_anno_color def augment_prompt_with_trigger_word( self, prompt: str, face_caption: str, key_parsing_mask_dict = None, facial_token = "<|facial|>", max_num_facials = 5, num_id_images: int = 1, device: Optional[torch.device] = None, ): device = device or self._execution_device # face_caption_align: 'The person has one nose <|facial|>, two ears <|facial|>, two eyes <|facial|>, and a mouth <|facial|>, ' face_caption_align, key_parsing_mask_dict_align = insert_markers_to_prompt(face_caption, key_parsing_mask_dict) prompt_face = prompt + " Detail: " + face_caption_align max_text_length=330 if len(self.tokenizer(prompt_face, max_length=self.tokenizer.model_max_length, padding="max_length", truncation=False, return_tensors="pt").input_ids[0]) != 77: # Put face_caption_align at the beginning of the prompt, so that the original prompt is truncated, # but the face_caption_align is well kept. prompt_face = "Detail: " + face_caption_align + " Caption:" + prompt # Remove "<|facial|>" from prompt_face. # augmented_prompt: 'A person, police officer, half body shot Detail: # The person has one nose , two ears , two eyes , and a mouth , ' augmented_prompt = prompt_face.replace("<|facial|>", "") tokenizer = self.tokenizer facial_token_id = tokenizer.convert_tokens_to_ids(facial_token) image_token_id = None # image_token_id: the token id of "<|image|>". Disabled, as it's set to None. # facial_token_id: the token id of "<|facial|>". clean_input_id, image_token_mask, facial_token_mask = \ tokenize_and_mask_noun_phrases_ends(prompt_face, image_token_id, facial_token_id, tokenizer) image_token_idx, image_token_idx_mask, facial_token_idx, facial_token_idx_mask = \ prepare_image_token_idx(image_token_mask, facial_token_mask, num_id_images, max_num_facials) return augmented_prompt, clean_input_id, key_parsing_mask_dict_align, facial_token_mask, facial_token_idx, facial_token_idx_mask @torch.inference_mode() def extract_parsed_image_parts(self, input_image_obj, key_parsing_mask_dict, image_size=512, max_num_facials=5): facial_masks = [] parsed_image_parts = [] key_masked_raw_images_dict = {} transform_mask = transforms.Compose([transforms.CenterCrop(size=image_size), transforms.ToTensor(),]) clip_preprocessor = CLIPImageProcessor() num_facial_part = len(key_parsing_mask_dict) for key in key_parsing_mask_dict: key_mask=key_parsing_mask_dict[key] facial_masks.append(transform_mask(key_mask)) key_masked_raw_image = apply_mask_to_raw_image(input_image_obj, key_mask) key_masked_raw_images_dict[key] = key_masked_raw_image # clip_preprocessor normalizes key_masked_raw_image, so that (masked) zero pixels become non-zero. # It also resizes the image to 224x224. parsed_image_part = clip_preprocessor(images=key_masked_raw_image, return_tensors="pt").pixel_values parsed_image_parts.append(parsed_image_part) padding_ficial_clip_image = torch.zeros_like(torch.zeros([1, 3, 224, 224])) padding_ficial_mask = torch.zeros_like(torch.zeros([1, image_size, image_size])) if num_facial_part < max_num_facials: parsed_image_parts += [ torch.zeros_like(padding_ficial_clip_image) for _ in range(max_num_facials - num_facial_part) ] facial_masks += [ torch.zeros_like(padding_ficial_mask) for _ in range(max_num_facials - num_facial_part) ] parsed_image_parts = torch.stack(parsed_image_parts, dim=1).squeeze(0) facial_masks = torch.stack(facial_masks, dim=0).squeeze(dim=1) return parsed_image_parts, facial_masks, key_masked_raw_images_dict # Release the unet and vae models to save memory. def release_unet_vae(self): unet = edict() # Only keep the config and in_channels attributes that are used in the pipeline. unet.config = self.unet.config unet.in_channels = self.unet.in_channels self.unet = unet self.vae = None # input_subj_image_obj: an Image object. def generate_id_prompt_embeds(self, prompt, negative_prompt, input_subj_image_obj, device, calc_uncond=True): faceid_embeds = self.extract_faceid(face_image=input_subj_image_obj) face_caption = "The person has one nose, two eyes, two ears, and a mouth." key_parsing_mask_dict, vis_parsing_anno_color = self.extract_facemask(input_subj_image_obj) augmented_prompt, clean_input_id, key_parsing_mask_dict_align, \ facial_token_mask, facial_token_idx, facial_token_idx_mask \ = self.augment_prompt_with_trigger_word( prompt = prompt, face_caption = face_caption, key_parsing_mask_dict=key_parsing_mask_dict, device=device, max_num_facials = 5, num_id_images = 1 ) text_embeds, uncond_text_embeds = self.encode_prompt( augmented_prompt, device=device, num_images_per_prompt=1, do_classifier_free_guidance=calc_uncond, negative_prompt=negative_prompt, ) # 5. Prepare the input ID images # global_id_embeds: [1, 4, 768] # get_global_id_embeds() extrats OpenCLIP embeddings from the input image and map them to global face prompt embeddings. global_id_embeds, uncond_global_id_embeds = \ self.get_global_id_embeds(faceid_embeds, face_image=input_subj_image_obj, s_scale=1.0, shortcut=False) # parsed_image_parts: [5, 3, 224, 224]. 5 parts, each part is a 3-channel 224x224 image (resized by CLIP Preprocessor). parsed_image_parts, facial_masks, key_masked_raw_images_dict = \ self.extract_parsed_image_parts(input_subj_image_obj, key_parsing_mask_dict_align, image_size=512, max_num_facials=5) parsed_image_parts2 = parsed_image_parts.unsqueeze(0).to(device, dtype=self.torch_dtype) facial_token_mask = facial_token_mask.to(device) facial_token_idx_mask = facial_token_idx_mask.to(device) # key_masked_raw_images_dict: ['Right_Eye', 'Right_Ear', 'Nose', 'Upper_Lip'] # for key in key_masked_raw_images_dict: # key_masked_raw_images_dict[key].save(f"{key}.png") # 6. Get the update text embedding # parsed_image_parts2: the facial areas of the input image # text_local_id_embeds: [1, 77, 768] # text_local_id_embeds only differs with text_global_id_embeds on 4 tokens, and is identical # to text_global_id_embeds on the rest 73 tokens. # get_local_facial_embeds() maps parsed_image_parts2 to multi_facial_embeds, and then replaces the class tokens in prompt_embeds # with the fused (id_embeds, prompt_embeds[class_tokens_mask]) whose indices are specified by class_tokens_mask. # parsed_image_parts2: [1, 5, 3, 224, 224] text_local_id_embeds, uncond_text_local_id_embeds = \ self.get_local_facial_embeds(text_embeds, uncond_text_embeds, \ parsed_image_parts2, facial_token_mask, facial_token_idx_mask, calc_uncond=calc_uncond) # text_global_id_embeds, text_local_global_id_embeds: [1, 81, 768] text_global_id_embeds = torch.cat([text_embeds, global_id_embeds], dim=1) text_local_global_id_embeds = torch.cat([text_local_id_embeds, global_id_embeds], dim=1) if calc_uncond: uncond_text_global_id_embeds = torch.cat([uncond_text_local_id_embeds, uncond_global_id_embeds], dim=1) coarse_prompt_embeds = torch.cat([uncond_text_global_id_embeds, text_global_id_embeds], dim=0) fine_prompt_embeds = torch.cat([uncond_text_global_id_embeds, text_local_global_id_embeds], dim=0) else: coarse_prompt_embeds = text_global_id_embeds fine_prompt_embeds = text_local_global_id_embeds # fine_prompt_embeds: the conditional part is # (text_global_id_embeds + text_local_global_id_embeds) / 2. fine_prompt_embeds = (coarse_prompt_embeds + fine_prompt_embeds) / 2 return coarse_prompt_embeds, fine_prompt_embeds @torch.no_grad() def __call__( self, prompt: Union[str, List[str]] = None, height: Optional[int] = None, width: Optional[int] = None, num_inference_steps: int = 50, guidance_scale: float = 5.0, negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, latents: Optional[torch.FloatTensor] = None, prompt_embeds: Optional[torch.FloatTensor] = None, negative_prompt_embeds: Optional[torch.FloatTensor] = None, output_type: Optional[str] = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, original_size: Optional[Tuple[int, int]] = None, target_size: Optional[Tuple[int, int]] = None, callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, callback_steps: int = 1, input_subj_image_objs: PipelineImageInput = None, start_merge_step: int = 0, ): # 0. Default height and width to unet height = height or self.unet.config.sample_size * self.vae_scale_factor width = width or self.unet.config.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( prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds, ) if not isinstance(input_subj_image_objs, list): input_subj_image_objs = [input_subj_image_objs] # 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] device = self._execution_device do_classifier_free_guidance = guidance_scale >= 1.0 assert do_classifier_free_guidance # 3. Encode input prompt coarse_prompt_embeds, fine_prompt_embeds = \ self.generate_id_prompt_embeds(prompt, negative_prompt, input_subj_image_objs[0], device) # 7. Prepare timesteps self.scheduler.set_timesteps(num_inference_steps, device=device) timesteps = self.scheduler.timesteps # 8. Prepare latent variables num_channels_latents = self.unet.in_channels latents = self.prepare_latents( batch_size * num_images_per_prompt, num_channels_latents, height, width, coarse_prompt_embeds.dtype, device, generator, latents, ) extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) cross_attention_kwargs = {} # 9. Denoising loop num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): latent_model_input = ( torch.cat([latents] * 2) if do_classifier_free_guidance else latents ) latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) if i <= start_merge_step: current_prompt_embeds = coarse_prompt_embeds else: current_prompt_embeds = fine_prompt_embeds # predict the noise residual noise_pred = self.unet( latent_model_input, t, encoder_hidden_states=current_prompt_embeds, cross_attention_kwargs=cross_attention_kwargs, ).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 ) else: assert 0, "Not Implemented" # 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 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: callback(i, t, latents) if output_type == "latent": image = latents elif output_type == "pil": # 9.1 Post-processing image = self.decode_latents(latents) # 9.3 Convert to PIL image = self.numpy_to_pil(image) else: # 9.1 Post-processing image = self.decode_latents(latents) # Offload last model to CPU if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: self.final_offload_hook.offload() if not return_dict: return (image, None) return StableDiffusionPipelineOutput( images=image, nsfw_content_detected=None )