# 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 math
import numpy
import torch
import inspect
import warnings
from PIL import Image
from einops import rearrange
import torch.nn.functional as F
from diffusers.utils.torch_utils import randn_tensor
from diffusers.configuration_utils import FrozenDict
from diffusers.image_processor import VaeImageProcessor
from typing import Any, Callable, Dict, List, Optional, Union
from diffusers.models import AutoencoderKL, UNet2DConditionModel
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers import DDPMScheduler, EulerAncestralDiscreteScheduler, ImagePipelineOutput
from diffusers.loaders import (
FromSingleFileMixin,
LoraLoaderMixin,
TextualInversionLoaderMixin
)
from transformers import (
CLIPImageProcessor,
CLIPTextModel,
CLIPTokenizer,
CLIPVisionModelWithProjection
)
from diffusers.models.attention_processor import (
Attention,
AttnProcessor,
XFormersAttnProcessor,
AttnProcessor2_0
)
from .utils import to_rgb_image, white_out_background, recenter_img
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> import torch
>>> from here import Hunyuan3d_MVD_Lite_Pipeline
>>> pipe = Hunyuan3d_MVD_Lite_Pipeline.from_pretrained(
... "weights/mvd_lite", torch_dtype=torch.float16
... )
>>> pipe.to("cuda")
>>> img = Image.open("demo.png")
>>> res_img = pipe(img).images[0]
"""
def unscale_latents(latents): return latents / 0.75 + 0.22
def unscale_image (image ): return image / 0.50 * 0.80
def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
return noise_cfg
class ReferenceOnlyAttnProc(torch.nn.Module):
# reference attention
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):
if encoder_hidden_states is None: encoder_hidden_states = 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)
res = self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask)
return res
class RefOnlyNoisedUNet(torch.nn.Module):
def __init__(self, unet, train_sched, val_sched):
super().__init__()
self.unet = unet
self.train_sched = train_sched
self.val_sched = val_sched
unet_lora_attn_procs = dict()
for name, _ in unet.attn_processors.items():
unet_lora_attn_procs[name] = ReferenceOnlyAttnProc(AttnProcessor2_0(),
enabled=name.endswith("attn1.processor"),
name=name)
unet.set_attn_processor(unet_lora_attn_procs)
def __getattr__(self, name: str):
try:
return super().__getattr__(name)
except AttributeError:
return getattr(self.unet, name)
def forward(self, sample, timestep, encoder_hidden_states, *args, cross_attention_kwargs, **kwargs):
cond_lat = cross_attention_kwargs['cond_lat']
noise = torch.randn_like(cond_lat)
if self.training:
noisy_cond_lat = self.train_sched.add_noise(cond_lat, noise, timestep)
noisy_cond_lat = self.train_sched.scale_model_input(noisy_cond_lat, timestep)
else:
noisy_cond_lat = self.val_sched.add_noise(cond_lat, noise, timestep.reshape(-1))
noisy_cond_lat = self.val_sched.scale_model_input(noisy_cond_lat, timestep.reshape(-1))
ref_dict = {}
self.unet(noisy_cond_lat,
timestep,
encoder_hidden_states,
*args,
cross_attention_kwargs=dict(mode="w", ref_dict=ref_dict),
**kwargs)
return self.unet(sample,
timestep,
encoder_hidden_states,
*args,
cross_attention_kwargs=dict(mode="r", ref_dict=ref_dict),
**kwargs)
class Hunyuan3d_MVD_Lite_Pipeline(DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin, FromSingleFileMixin):
def __init__(
self,
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
unet: UNet2DConditionModel,
scheduler: KarrasDiffusionSchedulers,
vision_encoder: CLIPVisionModelWithProjection,
feature_extractor_clip: CLIPImageProcessor,
feature_extractor_vae: CLIPImageProcessor,
ramping_coefficients: Optional[list] = None,
safety_checker=None,
):
DiffusionPipeline.__init__(self)
self.register_modules(
vae=vae,
unet=unet,
tokenizer=tokenizer,
scheduler=scheduler,
text_encoder=text_encoder,
vision_encoder=vision_encoder,
feature_extractor_vae=feature_extractor_vae,
feature_extractor_clip=feature_extractor_clip
)
# rewrite the stable diffusion pipeline
# vae: vae
# unet: unet
# tokenizer: tokenizer
# scheduler: scheduler
# text_encoder: text_encoder
# vision_encoder: vision_encoder
# feature_extractor_vae: feature_extractor_vae
# feature_extractor_clip: feature_extractor_clip
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)
def prepare_extra_step_kwargs(self, generator, eta):
extra_step_kwargs = {}
accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
if accepts_eta: extra_step_kwargs["eta"] = eta
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
if accepts_generator: extra_step_kwargs["generator"] = generator
return extra_step_kwargs
def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
latents = latents * self.scheduler.init_noise_sigma
return latents
@torch.no_grad()
def _encode_prompt(
self,
prompt,
device,
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
lora_scale: Optional[float] = None,
):
if lora_scale is not None and isinstance(self, LoraLoaderMixin):
self._lora_scale = lora_scale
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 prompt_embeds is None:
if isinstance(self, TextualInversionLoaderMixin):
prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
text_inputs = self.tokenizer(
prompt,
padding="max_length",
max_length=self.tokenizer.model_max_length,
truncation=True,
return_tensors="pt",
)
text_input_ids = text_inputs.input_ids
if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
attention_mask = text_inputs.attention_mask.to(device)
else:
attention_mask = None
prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)[0]
if self.text_encoder is not None:
prompt_embeds_dtype = self.text_encoder.dtype
elif self.unet is not None:
prompt_embeds_dtype = self.unet.dtype
else:
prompt_embeds_dtype = prompt_embeds.dtype
prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
bs_embed, seq_len, _ = prompt_embeds.shape
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
if do_classifier_free_guidance and negative_prompt_embeds is None:
uncond_tokens: List[str]
if negative_prompt is None: uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError()
elif isinstance(negative_prompt, str): uncond_tokens = [negative_prompt]
elif batch_size != len(negative_prompt): raise ValueError()
else: uncond_tokens = negative_prompt
if isinstance(self, TextualInversionLoaderMixin):
uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)
max_length = prompt_embeds.shape[1]
uncond_input = self.tokenizer(uncond_tokens,
padding="max_length",
max_length=max_length,
truncation=True,
return_tensors="pt")
if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
attention_mask = uncond_input.attention_mask.to(device)
else:
attention_mask = None
negative_prompt_embeds = self.text_encoder(uncond_input.input_ids.to(device), attention_mask=attention_mask)
negative_prompt_embeds = negative_prompt_embeds[0]
if do_classifier_free_guidance:
seq_len = negative_prompt_embeds.shape[1]
negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
return prompt_embeds
@torch.no_grad()
def encode_condition_image(self, image: torch.Tensor): return self.vae.encode(image).latent_dist.sample()
@torch.no_grad()
def __call__(self, image=None,
width=640,
height=960,
num_inference_steps=75,
return_dict=True,
generator=None,
**kwargs):
batch_size = 1
num_images_per_prompt = 1
output_type = 'pil'
do_classifier_free_guidance = True
guidance_rescale = 0.
if isinstance(self.unet, UNet2DConditionModel):
self.unet = RefOnlyNoisedUNet(self.unet, None, self.scheduler).eval()
cond_image = recenter_img(image)
cond_image = to_rgb_image(image)
image = cond_image
image_1 = self.feature_extractor_vae(images=image, return_tensors="pt").pixel_values
image_2 = self.feature_extractor_clip(images=image, return_tensors="pt").pixel_values
image_1 = image_1.to(device=self.vae.device, dtype=self.vae.dtype)
image_2 = image_2.to(device=self.vae.device, dtype=self.vae.dtype)
cond_lat = self.encode_condition_image(image_1)
negative_lat = self.encode_condition_image(torch.zeros_like(image_1))
cond_lat = torch.cat([negative_lat, cond_lat])
cross_attention_kwargs = dict(cond_lat=cond_lat)
global_embeds = self.vision_encoder(image_2, output_hidden_states=False).image_embeds.unsqueeze(-2)
encoder_hidden_states = self._encode_prompt('', self.device, num_images_per_prompt, False)
ramp = global_embeds.new_tensor(self.config.ramping_coefficients).unsqueeze(-1)
prompt_embeds = torch.cat([encoder_hidden_states, encoder_hidden_states + global_embeds * ramp])
device = self._execution_device
self.scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = self.scheduler.timesteps
num_channels_latents = self.unet.config.in_channels
latents = self.prepare_latents(batch_size * num_images_per_prompt,
num_channels_latents,
height,
width,
prompt_embeds.dtype,
device,
generator,
None)
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, 0.0)
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
# set adaptive cfg
# the image order is:
# [0, 60,
# 120, 180,
# 240, 300]
# the cfg is set as 3, 2.5, 2, 1.5
tmp_guidance_scale = torch.ones_like(latents)
tmp_guidance_scale[:, :, :40, :40] = 3
tmp_guidance_scale[:, :, :40, 40:] = 2.5
tmp_guidance_scale[:, :, 40:80, :40] = 2
tmp_guidance_scale[:, :, 40:80, 40:] = 1.5
tmp_guidance_scale[:, :, 80:120, :40] = 2
tmp_guidance_scale[:, :, 80:120, 40:] = 2.5
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)
noise_pred = self.unet(latent_model_input, t,
encoder_hidden_states=prompt_embeds,
cross_attention_kwargs=cross_attention_kwargs,
return_dict=False)[0]
adaptive_guidance_scale = (2 + 16 * (t / 1000) ** 5) / 3
if do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + \
tmp_guidance_scale * adaptive_guidance_scale * \
(noise_pred_text - noise_pred_uncond)
if do_classifier_free_guidance and guidance_rescale > 0.0:
noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
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)
image = unscale_image(self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0])
image = self.image_processor.postprocess(image, output_type='pil')[0]
image = [image, cond_image]
return ImagePipelineOutput(images=image) if return_dict else (image,)