Spaces:
Running
on
T4
Running
on
T4
File size: 3,856 Bytes
514160d |
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 |
# 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 random
import numpy as np
import torch
from typing import List
from diffusers import DiffusionPipeline
from diffusers import EulerAncestralDiscreteScheduler, LCMScheduler
class Multiview_Diffusion_Net():
def __init__(self, config) -> None:
self.device = config.device
self.view_size = 512
multiview_ckpt_path = config.multiview_ckpt_path
current_file_path = os.path.abspath(__file__)
custom_pipeline_path = os.path.join(os.path.dirname(current_file_path), '..', 'hunyuanpaint')
pipeline = DiffusionPipeline.from_pretrained(
multiview_ckpt_path,
custom_pipeline=custom_pipeline_path, torch_dtype=torch.float16)
if config.pipe_name in ['hunyuanpaint']:
pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config,
timestep_spacing='trailing')
elif config.pipe_name in ['hunyuanpaint-turbo']:
pipeline.scheduler = LCMScheduler.from_config(pipeline.scheduler.config,
timestep_spacing='trailing')
pipeline.set_turbo(True)
# pipeline.prepare()
pipeline.set_progress_bar_config(disable=True)
self.pipeline = pipeline.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)
def __call__(self, input_images, control_images, camera_info):
self.seed_everything(0)
if not isinstance(input_images, List):
input_images = [input_images]
input_images = [input_image.resize((self.view_size, self.view_size)) for input_image in input_images]
for i in range(len(control_images)):
control_images[i] = control_images[i].resize((self.view_size, self.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)]]
camera_info_gen = [camera_info]
camera_info_ref = [[0]]
kwargs['width'] = self.view_size
kwargs['height'] = self.view_size
kwargs['num_in_batch'] = num_view
kwargs['camera_info_gen'] = camera_info_gen
kwargs['camera_info_ref'] = camera_info_ref
kwargs["normal_imgs"] = normal_image
kwargs["position_imgs"] = position_image
mvd_image = self.pipeline(input_images, num_inference_steps=30, **kwargs).images
return mvd_image
|