File size: 5,536 Bytes
eb7ca4d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
# 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 os
import torch
import random
import numpy as np
from PIL import Image
from typing import List
import huggingface_hub
from omegaconf import OmegaConf
from diffusers import DiffusionPipeline
from diffusers import EulerAncestralDiscreteScheduler, DDIMScheduler, UniPCMultistepScheduler
class multiviewDiffusionNet:
def __init__(self, config) -> None:
self.device = config.device
cfg_path = config.multiview_cfg_path
custom_pipeline = config.custom_pipeline
cfg = OmegaConf.load(cfg_path)
self.cfg = cfg
self.mode = self.cfg.model.params.stable_diffusion_config.custom_pipeline[2:]
model_path = huggingface_hub.snapshot_download(
repo_id=config.multiview_pretrained_path,
allow_patterns=["hunyuan3d-paintpbr-v2-1/*"],
)
model_path = os.path.join(model_path, "hunyuan3d-paintpbr-v2-1")
pipeline = DiffusionPipeline.from_pretrained(
model_path,
custom_pipeline=custom_pipeline,
torch_dtype=torch.float16
)
pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config, timestep_spacing="trailing")
pipeline.set_progress_bar_config(disable=True)
pipeline.eval()
setattr(pipeline, "view_size", cfg.model.params.get("view_size", 320))
self.pipeline = pipeline.to(self.device)
if hasattr(self.pipeline.unet, "use_dino") and self.pipeline.unet.use_dino:
from hunyuanpaintpbr.unet.modules import Dino_v2
self.dino_v2 = Dino_v2(config.dino_ckpt_path).to(torch.float16)
self.dino_v2 = self.dino_v2.to(self.device)
def seed_everything(self, seed):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
os.environ["PL_GLOBAL_SEED"] = str(seed)
@torch.no_grad()
def __call__(self, images, conditions, prompt=None, custom_view_size=None, resize_input=False):
pils = self.forward_one(
images, conditions, prompt=prompt, custom_view_size=custom_view_size, resize_input=resize_input
)
return pils
def forward_one(self, input_images, control_images, prompt=None, custom_view_size=None, resize_input=False):
self.seed_everything(0)
custom_view_size = custom_view_size if custom_view_size is not None else self.pipeline.view_size
if not isinstance(input_images, List):
input_images = [input_images]
if not resize_input:
input_images = [
input_image.resize((self.pipeline.view_size, self.pipeline.view_size)) for input_image in input_images
]
else:
input_images = [input_image.resize((custom_view_size, custom_view_size)) for input_image in input_images]
for i in range(len(control_images)):
control_images[i] = control_images[i].resize((custom_view_size, custom_view_size))
if control_images[i].mode == "L":
control_images[i] = control_images[i].point(lambda x: 255 if x > 1 else 0, mode="1")
kwargs = dict(generator=torch.Generator(device=self.pipeline.device).manual_seed(0))
num_view = len(control_images) // 2
normal_image = [[control_images[i] for i in range(num_view)]]
position_image = [[control_images[i + num_view] for i in range(num_view)]]
kwargs["width"] = custom_view_size
kwargs["height"] = custom_view_size
kwargs["num_in_batch"] = num_view
kwargs["images_normal"] = normal_image
kwargs["images_position"] = position_image
if hasattr(self.pipeline.unet, "use_dino") and self.pipeline.unet.use_dino:
dino_hidden_states = self.dino_v2(input_images[0])
kwargs["dino_hidden_states"] = dino_hidden_states
sync_condition = None
infer_steps_dict = {
"EulerAncestralDiscreteScheduler": 30,
"UniPCMultistepScheduler": 15,
"DDIMScheduler": 50,
"ShiftSNRScheduler": 15,
}
mvd_image = self.pipeline(
input_images[0:1],
num_inference_steps=infer_steps_dict[self.pipeline.scheduler.__class__.__name__],
prompt=prompt,
sync_condition=sync_condition,
guidance_scale=3.0,
**kwargs,
).images
if "pbr" in self.mode:
mvd_image = {"albedo": mvd_image[:num_view], "mr": mvd_image[num_view:]}
# mvd_image = {'albedo':mvd_image[:num_view]}
else:
mvd_image = {"hdr": mvd_image}
return mvd_image
|