# Open Source Model Licensed under the Apache License Version 2.0 # and Other Licenses of the Third-Party Components therein: # The below Model in this distribution may have been modified by THL A29 Limited # ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited. # Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved. # The below software and/or models in this distribution may have been # modified by THL A29 Limited ("Tencent Modifications"). # All Tencent Modifications are Copyright (C) THL A29 Limited. # Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT # except for the third-party components listed below. # Hunyuan 3D does not impose any additional limitations beyond what is outlined # in the repsective licenses of these third-party components. # Users must comply with all terms and conditions of original licenses of these third-party # components and must ensure that the usage of the third party components adheres to # all relevant laws and regulations. # For avoidance of doubts, Hunyuan 3D means the large language models and # their software and algorithms, including trained model weights, parameters (including # optimizer states), machine-learning model code, inference-enabling code, training-enabling code, # fine-tuning enabling code and other elements of the foregoing made publicly available # by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT. import inspect from typing import Any, Dict, Optional from typing import Any, Dict, List, Optional, Tuple, Union import os import torch import numpy as np from PIL import Image import diffusers from diffusers.image_processor import VaeImageProcessor from diffusers.utils.import_utils import is_xformers_available from diffusers.schedulers import KarrasDiffusionSchedulers from diffusers.utils.torch_utils import randn_tensor from diffusers.utils.import_utils import is_xformers_available from diffusers.models.attention_processor import ( Attention, AttnProcessor, XFormersAttnProcessor, AttnProcessor2_0 ) from diffusers import ( AutoencoderKL, DDPMScheduler, DiffusionPipeline, EulerAncestralDiscreteScheduler, UNet2DConditionModel, ImagePipelineOutput ) import transformers from transformers import ( CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection, CLIPTextModelWithProjection ) from .utils import to_rgb_image, white_out_background, recenter_img EXAMPLE_DOC_STRING = """ Examples: ```py >>> import torch >>> from diffusers import Hunyuan3d_MVD_XL_Pipeline >>> pipe = Hunyuan3d_MVD_XL_Pipeline.from_pretrained( ... "Tencent-Hunyuan-3D/MVD-XL", torch_dtype=torch.float16 ... ) >>> pipe.to("cuda") >>> img = Image.open("demo.png") >>> res_img = pipe(img).images[0] ``` """ def scale_latents(latents): return (latents - 0.22) * 0.75 def unscale_latents(latents): return (latents / 0.75) + 0.22 def scale_image(image): return (image - 0.5) / 0.5 def scale_image_2(image): return (image * 0.5) / 0.8 def unscale_image(image): return (image * 0.5) + 0.5 def unscale_image_2(image): return (image * 0.8) / 0.5 class ReferenceOnlyAttnProc(torch.nn.Module): def __init__(self, chained_proc, enabled=False, name=None): super().__init__() self.enabled = enabled self.chained_proc = chained_proc self.name = name def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, mode="w", ref_dict=None): encoder_hidden_states = hidden_states if encoder_hidden_states is None else encoder_hidden_states if self.enabled: if mode == 'w': ref_dict[self.name] = encoder_hidden_states elif mode == 'r': encoder_hidden_states = torch.cat([encoder_hidden_states, ref_dict.pop(self.name)], dim=1) else: raise Exception(f"mode should not be {mode}") return self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask) class RefOnlyNoisedUNet(torch.nn.Module): def __init__(self, unet, scheduler) -> None: super().__init__() self.unet = unet self.scheduler = scheduler unet_attn_procs = dict() for name, _ in unet.attn_processors.items(): if torch.__version__ >= '2.0': default_attn_proc = AttnProcessor2_0() elif is_xformers_available(): default_attn_proc = XFormersAttnProcessor() else: default_attn_proc = AttnProcessor() unet_attn_procs[name] = ReferenceOnlyAttnProc( default_attn_proc, enabled=name.endswith("attn1.processor"), name=name ) unet.set_attn_processor(unet_attn_procs) def __getattr__(self, name: str): try: return super().__getattr__(name) except AttributeError: return getattr(self.unet, name) def forward( self, sample: torch.FloatTensor, timestep: Union[torch.Tensor, float, int], encoder_hidden_states: torch.Tensor, cross_attention_kwargs: Optional[Dict[str, Any]] = None, class_labels: Optional[torch.Tensor] = None, down_block_res_samples: Optional[Tuple[torch.Tensor]] = None, mid_block_res_sample: Optional[Tuple[torch.Tensor]] = None, added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, return_dict: bool = True, **kwargs ): dtype = self.unet.dtype # cond_lat add same level noise cond_lat = cross_attention_kwargs['cond_lat'] noise = torch.randn_like(cond_lat) noisy_cond_lat = self.scheduler.add_noise(cond_lat, noise, timestep.reshape(-1)) noisy_cond_lat = self.scheduler.scale_model_input(noisy_cond_lat, timestep.reshape(-1)) ref_dict = {} _ = self.unet( noisy_cond_lat, timestep, encoder_hidden_states = encoder_hidden_states, class_labels = class_labels, cross_attention_kwargs = dict(mode="w", ref_dict=ref_dict), added_cond_kwargs = added_cond_kwargs, return_dict = return_dict, **kwargs ) res = self.unet( sample, timestep, encoder_hidden_states, class_labels=class_labels, cross_attention_kwargs = dict(mode="r", ref_dict=ref_dict), down_block_additional_residuals = [ sample.to(dtype=dtype) for sample in down_block_res_samples ] if down_block_res_samples is not None else None, mid_block_additional_residual = ( mid_block_res_sample.to(dtype=dtype) if mid_block_res_sample is not None else None), added_cond_kwargs = added_cond_kwargs, return_dict = return_dict, **kwargs ) return res class HunYuan3D_MVD_Std_Pipeline(diffusers.DiffusionPipeline): def __init__( self, vae: AutoencoderKL, unet: UNet2DConditionModel, scheduler: KarrasDiffusionSchedulers, feature_extractor_vae: CLIPImageProcessor, vision_processor: CLIPImageProcessor, vision_encoder: CLIPVisionModelWithProjection, vision_encoder_2: CLIPVisionModelWithProjection, ramping_coefficients: Optional[list] = None, add_watermarker: Optional[bool] = None, safety_checker = None, ): DiffusionPipeline.__init__(self) self.register_modules( vae=vae, unet=unet, scheduler=scheduler, safety_checker=None, feature_extractor_vae=feature_extractor_vae, vision_processor=vision_processor, vision_encoder=vision_encoder, vision_encoder_2=vision_encoder_2, ) self.register_to_config( ramping_coefficients = ramping_coefficients) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) self.default_sample_size = self.unet.config.sample_size self.watermark = None self.prepare_init = False def prepare(self): assert isinstance(self.unet, UNet2DConditionModel), "unet should be UNet2DConditionModel" self.unet = RefOnlyNoisedUNet(self.unet, self.scheduler).eval() self.prepare_init = True def encode_image(self, image: torch.Tensor, scale_factor: bool = False): latent = self.vae.encode(image).latent_dist.sample() return (latent * self.vae.config.scaling_factor) if scale_factor else latent # 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 _get_add_time_ids( self, original_size, crops_coords_top_left, target_size, dtype, text_encoder_projection_dim=None ): add_time_ids = list(original_size + crops_coords_top_left + target_size) passed_add_embed_dim = ( self.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim ) expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features if expected_add_embed_dim != passed_add_embed_dim: raise ValueError( f"Model expects an added time embedding vector of length {expected_add_embed_dim}, " \ f"but a vector of {passed_add_embed_dim} was created. The model has an incorrect config." \ f" Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`." ) add_time_ids = torch.tensor([add_time_ids], dtype=dtype) return add_time_ids 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 @property def guidance_scale(self): return self._guidance_scale @property def interrupt(self): return self._interrupt @property def do_classifier_free_guidance(self): return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None @torch.no_grad() def __call__( self, image: Image.Image = None, guidance_scale = 2.0, output_type: Optional[str] = "pil", num_inference_steps: int = 50, return_dict: bool = True, eta: float = 0.0, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, crops_coords_top_left: Tuple[int, int] = (0, 0), cross_attention_kwargs: Optional[Dict[str, Any]] = None, latent: torch.Tensor = None, guidance_curve = None, **kwargs ): if not self.prepare_init: self.prepare() here = dict(device=self.vae.device, dtype=self.vae.dtype) batch_size = 1 num_images_per_prompt = 1 width, height = 512 * 2, 512 * 3 target_size = original_size = (height, width) self._guidance_scale = guidance_scale self._cross_attention_kwargs = cross_attention_kwargs self._interrupt = False device = self._execution_device # Prepare timesteps self.scheduler.set_timesteps(num_inference_steps, device=device) timesteps = self.scheduler.timesteps # 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, self.vae.dtype, device, generator, latents=latent, ) # Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) # Prepare added time ids & embeddings text_encoder_projection_dim = 1280 add_time_ids = self._get_add_time_ids( original_size, crops_coords_top_left, target_size, dtype=self.vae.dtype, text_encoder_projection_dim=text_encoder_projection_dim, ) negative_add_time_ids = add_time_ids # hw: preprocess cond_image = recenter_img(image) cond_image = to_rgb_image(image) image_vae = self.feature_extractor_vae(images=cond_image, return_tensors="pt").pixel_values.to(**here) image_clip = self.vision_processor(images=cond_image, return_tensors="pt").pixel_values.to(**here) # hw: get cond_lat from cond_img using vae cond_lat = self.encode_image(image_vae, scale_factor=False) negative_lat = self.encode_image(torch.zeros_like(image_vae), scale_factor=False) cond_lat = torch.cat([negative_lat, cond_lat]) # hw: get visual global embedding using clip global_embeds_1 = self.vision_encoder(image_clip, output_hidden_states=False).image_embeds.unsqueeze(-2) global_embeds_2 = self.vision_encoder_2(image_clip, output_hidden_states=False).image_embeds.unsqueeze(-2) global_embeds = torch.concat([global_embeds_1, global_embeds_2], dim=-1) ramp = global_embeds.new_tensor(self.config.ramping_coefficients).unsqueeze(-1) prompt_embeds = self.uc_text_emb.to(**here) pooled_prompt_embeds = self.uc_text_emb_2.to(**here) prompt_embeds = prompt_embeds + global_embeds * ramp add_text_embeds = pooled_prompt_embeds if self.do_classifier_free_guidance: negative_prompt_embeds = torch.zeros_like(prompt_embeds) negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds) prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0) add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0) prompt_embeds = prompt_embeds.to(device) add_text_embeds = add_text_embeds.to(device) add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1) # Denoising loop num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) timestep_cond = None self._num_timesteps = len(timesteps) if guidance_curve is None: guidance_curve = lambda t: guidance_scale with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): if self.interrupt: continue # expand the latents if we are doing classifier free guidance latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) # predict the noise residual added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} noise_pred = self.unet( latent_model_input, t, encoder_hidden_states=prompt_embeds, timestep_cond=timestep_cond, cross_attention_kwargs=dict(cond_lat=cond_lat), added_cond_kwargs=added_cond_kwargs, return_dict=False, )[0] # perform guidance # cur_guidance_scale = self.guidance_scale cur_guidance_scale = guidance_curve(t) # 1.5 + 2.5 * ((t/1000)**2) if self.do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + cur_guidance_scale * (noise_pred_text - noise_pred_uncond) # cur_guidance_scale_topleft = (cur_guidance_scale - 1.0) * 4 + 1.0 # noise_pred_top_left = noise_pred_uncond + # cur_guidance_scale_topleft * (noise_pred_text - noise_pred_uncond) # _, _, h, w = noise_pred.shape # noise_pred[:, :, :h//3, :w//2] = noise_pred_top_left[:, :, :h//3, :w//2] # compute the previous noisy sample x_t -> x_t-1 latents_dtype = latents.dtype latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] # 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() latents = unscale_latents(latents) if output_type=="latent": image = latents else: image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] image = unscale_image(unscale_image_2(image)).clamp(0, 1) image = [ Image.fromarray((image[0]*255+0.5).clamp_(0, 255).permute(1, 2, 0).cpu().numpy().astype("uint8")), # self.image_processor.postprocess(image, output_type=output_type)[0], cond_image.resize((512, 512)) ] if not return_dict: return (image,) return ImagePipelineOutput(images=image) def save_pretrained(self, save_directory): # uc_text_emb.pt and uc_text_emb_2.pt are inferenced and saved in advance super().save_pretrained(save_directory) torch.save(self.uc_text_emb, os.path.join(save_directory, "uc_text_emb.pt")) torch.save(self.uc_text_emb_2, os.path.join(save_directory, "uc_text_emb_2.pt")) @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): # uc_text_emb.pt and uc_text_emb_2.pt are inferenced and saved in advance pipeline = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) pipeline.uc_text_emb = torch.load(os.path.join(pretrained_model_name_or_path, "uc_text_emb.pt")) pipeline.uc_text_emb_2 = torch.load(os.path.join(pretrained_model_name_or_path, "uc_text_emb_2.pt")) return pipeline