# Copyright 2023 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os from typing import List, Optional, Union, Tuple import numpy as np import torch from PIL import Image from transformers import ( AutoTokenizer, CLIPConfig, CLIPTextModel, CLIPTextModelWithProjection, PreTrainedTokenizer, T5EncoderModel, ) from .modeling_maskgit_vqgan import MaskGitVQGAN from .modeling_movq import MOVQ from .modeling_paella_vq import PaellaVQModel from .modeling_taming_vqgan import VQGANModel from .modeling_transformer import MaskGitTransformer, MaskGiTUViT from .sampling import get_mask_chedule class PipelineMuse: def __init__( self, vae: Union[VQGANModel, MOVQ, MaskGitVQGAN], transformer: Union[MaskGitTransformer, MaskGiTUViT], is_class_conditioned: bool = False, text_encoder: Optional[Union[T5EncoderModel, CLIPTextModel]] = None, tokenizer: Optional[PreTrainedTokenizer] = None, ) -> None: self.text_encoder = text_encoder self.tokenizer = tokenizer self.vae = vae self.transformer = transformer self.is_class_conditioned = is_class_conditioned self.device = "cpu" def to(self, device="cpu", dtype=torch.float32) -> None: self.device = device self.dtype = dtype if not self.is_class_conditioned: self.text_encoder.to(device, dtype=dtype) self.transformer.to(device, dtype=dtype) self.vae.to(device, dtype=torch.float32) # keep vae in fp32 return self @torch.no_grad() def __call__( self, text: Optional[Union[str, List[str]]] = None, negative_text: Optional[Union[str, List[str]]] = "", prompt_embeds: Optional[torch.Tensor] = None, pooled_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, negative_pooled_embeds: Optional[torch.Tensor] = None, class_ids: Optional[Union[int, List[int]]] = None, timesteps: int = 16, noise_schedule: str = "cosine", guidance_scale: float = 10.0, guidance_schedule=None, temperature: Union[float, Tuple[float]] = (2, 0), topk_filter_thres: float = 0.9, num_images_per_prompt: int = 1, use_maskgit_generate: bool = True, generator: Optional[torch.Generator] = None, use_fp16: bool = False, noise_type="mask", # can be "mask" or "random_replace" predict_all_tokens=False, orig_size=(512, 512), crop_coords=(0, 0), aesthetic_score=6.0, return_intermediate: bool = False, use_tqdm=True, transformer_seq_len=None, clip_skip:int = None, ): if text is None and class_ids is None: raise ValueError("Either text or class_ids must be provided.") if text is not None and class_ids is not None: raise ValueError("Only one of text or class_ids may be provided.") if class_ids is not None: if isinstance(class_ids, int): class_ids = [class_ids] class_ids = torch.tensor(class_ids, device=self.device, dtype=torch.long) # duplicate class ids for each generation per prompt class_ids = class_ids.repeat_interleave(num_images_per_prompt, dim=0) model_inputs = {"class_ids": class_ids} else: if isinstance(text, str): text = [text] if prompt_embeds is None: input_ids = self.tokenizer( text, return_tensors="pt", padding="max_length", truncation=True, max_length=self.tokenizer.model_max_length, ).input_ids # TODO: remove hardcode input_ids = input_ids.to(self.device) if self.transformer.config.add_cond_embeds: if prompt_embeds is not None and pooled_embeds is not None: pooled_embeds, encoder_hidden_states = pooled_embeds, prompt_embeds pooled_embeds = pooled_embeds.to(self.device, dtype=self.text_encoder.dtype) encoder_hidden_states = encoder_hidden_states.to(self.device, dtype=self.text_encoder.dtype) else: clip_layer_idx = -(clip_skip+1) if clip_skip is not None else -2 outputs = self.text_encoder(input_ids, return_dict=True, output_hidden_states=True) pooled_embeds, encoder_hidden_states = outputs.text_embeds, outputs.hidden_states[clip_layer_idx] else: encoder_hidden_states = self.text_encoder(input_ids).last_hidden_state pooled_embeds = None if negative_text is not None: if isinstance(negative_text, str): negative_text = [negative_text] * len(text) negative_input_ids = self.tokenizer( negative_text, return_tensors="pt", padding="max_length", truncation=True, max_length=self.tokenizer.model_max_length, ).input_ids negative_input_ids = negative_input_ids.to(self.device) if self.transformer.config.add_cond_embeds: outputs = self.text_encoder(negative_input_ids, return_dict=True, output_hidden_states=True) negative_pooled_embeds = outputs.text_embeds negative_encoder_hidden_states = outputs.hidden_states[-2] else: negative_encoder_hidden_states = self.text_encoder(negative_input_ids).last_hidden_state negative_pooled_embeds = None elif negative_prompt_embeds is not None: negative_encoder_hidden_states = negative_prompt_embeds.to(self.device, dtype=self.text_encoder.dtype) negative_pooled_embeds = negative_pooled_embeds.to(self.device, dtype=self.text_encoder.dtype) else: negative_encoder_hidden_states = None negative_pooled_embeds = None # duplicate text embeddings for each generation per prompt, using mps friendly method bs_embed, seq_len, _ = encoder_hidden_states.shape encoder_hidden_states = encoder_hidden_states.repeat(1, num_images_per_prompt, 1) encoder_hidden_states = encoder_hidden_states.view(bs_embed * num_images_per_prompt, seq_len, -1) if pooled_embeds is not None: bs_embed, _ = pooled_embeds.shape pooled_embeds = pooled_embeds.repeat(1, num_images_per_prompt) pooled_embeds = pooled_embeds.view(bs_embed * num_images_per_prompt, -1) if negative_pooled_embeds is not None: bs_embed, _ = negative_pooled_embeds.shape negative_pooled_embeds = negative_pooled_embeds.repeat(1, num_images_per_prompt) negative_pooled_embeds = negative_pooled_embeds.view(bs_embed * num_images_per_prompt, -1) if negative_encoder_hidden_states is not None: bs_embed, seq_len, _ = negative_encoder_hidden_states.shape negative_encoder_hidden_states = negative_encoder_hidden_states.repeat(1, num_images_per_prompt, 1) negative_encoder_hidden_states = negative_encoder_hidden_states.view( bs_embed * num_images_per_prompt, seq_len, -1 ) if negative_encoder_hidden_states is None: empty_input = self.tokenizer("", padding="max_length", return_tensors="pt").input_ids.to( self.text_encoder.device ) outputs = self.text_encoder(empty_input, output_hidden_states=True) empty_embeds = outputs.hidden_states[-2] empty_cond_embeds = outputs[0] else: empty_embeds, empty_cond_embeds = None, None model_inputs = { "encoder_hidden_states": encoder_hidden_states, "negative_embeds": negative_encoder_hidden_states, "cond_embeds": pooled_embeds, "negative_cond_embeds": negative_pooled_embeds, "empty_embeds": empty_embeds, "empty_cond_embeds": empty_cond_embeds, } if self.transformer.config.add_micro_cond_embeds: micro_conds = list(orig_size) + list(crop_coords) + [aesthetic_score] micro_conds = torch.tensor(micro_conds, device=self.device, dtype=encoder_hidden_states.dtype) micro_conds = micro_conds.unsqueeze(0) model_inputs["micro_conds"] = micro_conds generate = self.transformer.generate if use_maskgit_generate: generate = self.transformer.generate2 with torch.autocast("cuda", enabled=use_fp16): outputs = generate( **model_inputs, timesteps=timesteps, guidance_scale=guidance_scale, guidance_schedule=guidance_schedule, temperature=temperature, topk_filter_thres=topk_filter_thres, generator=generator, noise_type=noise_type, noise_schedule=get_mask_chedule(noise_schedule), predict_all_tokens=predict_all_tokens, return_intermediate=return_intermediate, use_tqdm=use_tqdm, seq_len=transformer_seq_len, ) if return_intermediate: generated_tokens, intermediate = outputs else: generated_tokens = outputs images = self.vae.decode_code(generated_tokens) if return_intermediate: intermediate_images = [self.vae.decode_code(tokens) for tokens in intermediate] # Convert to PIL images images = [self.to_pil_image(image) for image in images] if return_intermediate: intermediate_images = [[self.to_pil_image(image) for image in images] for images in intermediate_images] return images, intermediate_images return images def to_pil_image(self, image: torch.Tensor): image = image.permute(1, 2, 0).cpu().numpy() image = 2.0 * image - 1.0 image = np.clip(image, -1.0, 1.0) image = (image + 1.0) / 2.0 image = (255 * image).astype(np.uint8) image = Image.fromarray(image).convert("RGB") return image @classmethod def from_pretrained( cls, model_name_or_path: str = None, text_encoder_path: Optional[str] = None, vae_path: Optional[str] = None, transformer_path: Optional[str] = None, vae = None, text_encoder = None, transformer = None, is_class_conditioned: bool = False, ) -> None: """ Instantiate a PipelineMuse from a pretrained model. Either model_name_or_path or all of text_encoder_path, vae_path, and transformer_path must be provided. """ if model_name_or_path is None: if text_encoder_path is None or vae_path is None or transformer_path is None: raise ValueError( "If model_name_or_path is None, then text_encoder_path, vae_path, and transformer_path must be" " provided." ) text_encoder_args = None tokenizer_args = None if not is_class_conditioned: text_encoder_args = {"pretrained_model_name_or_path": text_encoder_path} tokenizer_args = {"pretrained_model_name_or_path": text_encoder_path} vae_args = {"pretrained_model_name_or_path": vae_path} transformer_args = {"pretrained_model_name_or_path": transformer_path} else: text_encoder_args = None tokenizer_args = None if not is_class_conditioned: text_encoder_args = {"pretrained_model_name_or_path": model_name_or_path, "subfolder": "text_encoder"} tokenizer_args = {"pretrained_model_name_or_path": model_name_or_path, "subfolder": "text_encoder"} vae_args = {"pretrained_model_name_or_path": model_name_or_path, "subfolder": "vae"} transformer_args = {"pretrained_model_name_or_path": model_name_or_path, "subfolder": "transformer"} if not is_class_conditioned: # Very hacky way to load different text encoders # TODO: Add config for pipeline to specify text encoder # is_clip = "clip" in text_encoder_args["pretrained_model_name_or_path"].lower() # text_encoder_cls = CLIPTextModel if is_clip else T5EncoderModel # if is_clip: # config = CLIPConfig.from_pretrained(**text_encoder_args) # if config.architectures[0] == "CLIPTextModel": # text_encoder_cls = CLIPTextModel # else: # text_encoder_cls = CLIPTextModelWithProjection # text_encoder_args["projection_dim"] = 768 # TODO: make this more robust if text_encoder is None: text_encoder = CLIPTextModelWithProjection.from_pretrained(**text_encoder_args) tokenizer = AutoTokenizer.from_pretrained(**tokenizer_args) transformer_config = MaskGitTransformer.load_config(**transformer_args) if transformer is not None: ... elif transformer_config["_class_name"] == "MaskGitTransformer": transformer = MaskGitTransformer.from_pretrained(**transformer_args) elif transformer_config["_class_name"] == "MaskGiTUViT" or transformer_config["_class_name"] == "MaskGiTUViT_v2": transformer = MaskGiTUViT.from_pretrained(**transformer_args) else: raise ValueError(f"Unknown Transformer class: {transformer_config['_class_name']}") # Hacky way to load different VQ models vae_config = MaskGitVQGAN.load_config(**vae_args) if vae is not None: ... elif vae_config["_class_name"] == "VQGANModel": vae = VQGANModel.from_pretrained(**vae_args) elif vae_config["_class_name"] == "MaskGitVQGAN": vae = MaskGitVQGAN.from_pretrained(**vae_args) elif vae_config["_class_name"] == "MOVQ": vae = MOVQ.from_pretrained(**vae_args) elif vae_config["_class_name"] == "PaellaVQModel": vae = PaellaVQModel.from_pretrained(**vae_args) else: raise ValueError(f"Unknown VAE class: {vae_config['_class_name']}") if is_class_conditioned: return cls( vae=vae, transformer=transformer, is_class_conditioned=is_class_conditioned, ) return cls( vae=vae, transformer=transformer, text_encoder=text_encoder, tokenizer=tokenizer, is_class_conditioned=is_class_conditioned, ) def save_pretrained( self, save_directory: Union[str, os.PathLike], ) -> None: """ Save the pipeline's model and tokenizer to the specified directory. """ if not self.is_class_conditioned: self.text_encoder.save_pretrained(os.path.join(save_directory, "text_encoder")) self.tokenizer.save_pretrained(os.path.join(save_directory, "text_encoder")) self.vae.save_pretrained(os.path.join(save_directory, "vae")) self.transformer.save_pretrained(os.path.join(save_directory, "transformer")) class PipelineMuseInpainting(PipelineMuse): @torch.no_grad() def __call__( self, image: Image, mask: torch.BoolTensor, text: Optional[Union[str, List[str]]] = None, negative_text: Optional[Union[str, List[str]]] = None, class_ids: torch.LongTensor = None, timesteps: int = 8, guidance_scale: float = 8.0, guidance_schedule=None, temperature: float = 1.0, topk_filter_thres: float = 0.9, num_images_per_prompt: int = 1, use_maskgit_generate: bool = True, generator: Optional[torch.Generator] = None, use_fp16: bool = False, image_size: int = 256, orig_size=(256, 256), crop_coords=(0, 0), aesthetic_score=6.0, ): from torchvision import transforms assert use_maskgit_generate if text is None and class_ids is None: raise ValueError("Either text or class_ids must be provided.") if text is not None and class_ids is not None: raise ValueError("Only one of text or class_ids may be provided.") encode_transform = transforms.Compose( [ transforms.Resize(image_size, interpolation=transforms.InterpolationMode.BILINEAR), transforms.CenterCrop(image_size), transforms.ToTensor(), ] ) pixel_values = encode_transform(image).unsqueeze(0).to(self.device) _, image_tokens = self.vae.encode(pixel_values) mask_token_id = self.transformer.config.mask_token_id image_tokens[mask[None]] = mask_token_id image_tokens = image_tokens.repeat(num_images_per_prompt, 1) if class_ids is not None: if isinstance(class_ids, int): class_ids = [class_ids] class_ids = torch.tensor(class_ids, device=self.device, dtype=torch.long) # duplicate class ids for each generation per prompt class_ids = class_ids.repeat_interleave(num_images_per_prompt, dim=0) model_inputs = {"class_ids": class_ids} else: if isinstance(text, str): text = [text] input_ids = self.tokenizer( text, return_tensors="pt", padding="max_length", truncation=True, max_length=self.tokenizer.model_max_length, ).input_ids # TODO: remove hardcode input_ids = input_ids.to(self.device) if self.transformer.config.add_cond_embeds: outputs = self.text_encoder(input_ids, return_dict=True, output_hidden_states=True) pooled_embeds, encoder_hidden_states = outputs.text_embeds, outputs.hidden_states[-2] else: encoder_hidden_states = self.text_encoder(input_ids).last_hidden_state pooled_embeds = None if negative_text is not None: if isinstance(negative_text, str): negative_text = [negative_text] negative_input_ids = self.tokenizer( negative_text, return_tensors="pt", padding="max_length", truncation=True, max_length=self.tokenizer.model_max_length, ).input_ids negative_input_ids = negative_input_ids.to(self.device) negative_encoder_hidden_states = self.text_encoder(negative_input_ids).last_hidden_state else: negative_encoder_hidden_states = None # duplicate text embeddings for each generation per prompt, using mps friendly method bs_embed, seq_len, _ = encoder_hidden_states.shape encoder_hidden_states = encoder_hidden_states.repeat(1, num_images_per_prompt, 1) encoder_hidden_states = encoder_hidden_states.view(bs_embed * num_images_per_prompt, seq_len, -1) if negative_encoder_hidden_states is not None: bs_embed, seq_len, _ = negative_encoder_hidden_states.shape negative_encoder_hidden_states = negative_encoder_hidden_states.repeat(1, num_images_per_prompt, 1) negative_encoder_hidden_states = negative_encoder_hidden_states.view( bs_embed * num_images_per_prompt, seq_len, -1 ) empty_input = self.tokenizer("", padding="max_length", return_tensors="pt").input_ids.to( self.text_encoder.device ) outputs = self.text_encoder(empty_input, output_hidden_states=True) empty_embeds = outputs.hidden_states[-2] empty_cond_embeds = outputs[0] model_inputs = { "encoder_hidden_states": encoder_hidden_states, "negative_embeds": negative_encoder_hidden_states, "empty_embeds": empty_embeds, "empty_cond_embeds": empty_cond_embeds, "cond_embeds": pooled_embeds, } if self.transformer.config.add_micro_cond_embeds: micro_conds = list(orig_size) + list(crop_coords) + [aesthetic_score] micro_conds = torch.tensor(micro_conds, device=self.device, dtype=encoder_hidden_states.dtype) micro_conds = micro_conds.unsqueeze(0) model_inputs["micro_conds"] = micro_conds generate = self.transformer.generate2 with torch.autocast("cuda", enabled=use_fp16): generated_tokens = generate( input_ids=image_tokens, **model_inputs, timesteps=timesteps, guidance_scale=guidance_scale, guidance_schedule=guidance_schedule, temperature=temperature, topk_filter_thres=topk_filter_thres, generator=generator, ) images = self.vae.decode_code(generated_tokens) # Convert to PIL images images = [self.to_pil_image(image) for image in images] return images