mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-08-26 08:19:07 +08:00
156 lines
5.3 KiB
Python
156 lines
5.3 KiB
Python
import torch
|
|
import trimesh
|
|
import torch.nn as nn
|
|
from PIL import Image
|
|
from typing import List, Union
|
|
from torch.utils._pytree import tree_map
|
|
from torch.utils.data._utils.collate import default_collate
|
|
|
|
def export_to_trimesh(mesh_output):
|
|
if isinstance(mesh_output, list):
|
|
outputs = []
|
|
for mesh in mesh_output:
|
|
if mesh is None:
|
|
outputs.append(None)
|
|
else:
|
|
mesh.mesh_f = mesh.mesh_f[:, ::-1]
|
|
mesh_output = trimesh.Trimesh(mesh.mesh_v, mesh.mesh_f)
|
|
outputs.append(mesh_output)
|
|
return outputs
|
|
else:
|
|
mesh_output.mesh_f = mesh_output.mesh_f[:, ::-1]
|
|
mesh_output = trimesh.Trimesh(mesh_output.mesh_v, mesh_output.mesh_f)
|
|
return mesh_output
|
|
|
|
class Hunyuan3DDiTFlowMatchingPipeline(nn.Module):
|
|
def __init__(self, model, vae, conditioner, image_processor, scheduler):
|
|
|
|
self.vae = vae
|
|
self.model = model
|
|
self.conditioner = conditioner
|
|
self.image_processor = image_processor
|
|
self.scheduler = scheduler
|
|
|
|
def compile(self):
|
|
self.vae = torch.compile(self.vae)
|
|
self.model = torch.compile(self.model)
|
|
self.conditioner = torch.compile(self.conditioner)
|
|
|
|
def encode_cond(self, image, additional_cond_inputs, do_classifier_free_guidance):
|
|
|
|
bsz = image.shape[0]
|
|
cond = self.conditioner(image=image, **additional_cond_inputs)
|
|
|
|
if do_classifier_free_guidance:
|
|
un_cond = self.conditioner.unconditional_embedding(bsz, **additional_cond_inputs)
|
|
|
|
# avoid python recursion by using tree_map
|
|
_fn = lambda x, y: torch.cat([x, y], dim=0).to(self.dtype)
|
|
|
|
cond = tree_map(_fn, cond, un_cond)
|
|
|
|
return cond
|
|
|
|
def prepare_images(self, images):
|
|
|
|
outputs = []
|
|
for image in images:
|
|
output = self.image_processor(image)
|
|
outputs.append(output)
|
|
|
|
return default_collate(outputs)
|
|
|
|
def get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=torch.float32):
|
|
|
|
assert len(w.shape) == 1
|
|
w = w * 1000.0
|
|
|
|
half_dim = embedding_dim // 2
|
|
emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
|
|
|
|
emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
|
|
emb = w.to(dtype)[:, None] * emb[None, :]
|
|
|
|
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
|
|
if embedding_dim % 2 == 1: # zero pad
|
|
emb = torch.nn.functional.pad(emb, (0, 1))
|
|
|
|
assert emb.shape == (w.shape[0], embedding_dim)
|
|
|
|
return emb
|
|
|
|
def prepare_latents(self, batch_size, dtype, device):
|
|
|
|
shape = (batch_size, *self.vae.latent_shape)
|
|
latents = torch.randn(shape, dtype = dtype, device = device)
|
|
|
|
return latents
|
|
|
|
@torch.inference_mode()
|
|
def __call__(
|
|
self,
|
|
image: Union[str, List[str], Image.Image, dict, List[dict], torch.Tensor] = None,
|
|
guidance_scale: float = 5.0,
|
|
bounds = 1.01,
|
|
octree_res = 384,
|
|
num_chunks = 8000,
|
|
**kwargs,
|
|
) -> List[List[trimesh.Trimesh]]:
|
|
|
|
callback = kwargs.pop("callback", None)
|
|
callback_steps = kwargs.pop("callback_steps", None)
|
|
|
|
device = self.device
|
|
dtype = self.dtype
|
|
do_classifier_free_guidance = guidance_scale >= 0 and not (
|
|
hasattr(self.model, 'guidance_embed') and
|
|
self.model.guidance_embed is True
|
|
)
|
|
|
|
cond_inputs = self.prepare_image(image)
|
|
image = cond_inputs.pop('image')
|
|
|
|
cond = self.encode_cond(
|
|
image = image,
|
|
additional_cond_inputs = cond_inputs,
|
|
do_classifier_free_guidance = do_classifier_free_guidance,
|
|
)
|
|
|
|
guidance = None
|
|
batch_size = image.shape[0]
|
|
|
|
latents = self.prepare_latents(batch_size, dtype, device)
|
|
|
|
if hasattr(self.model, 'guidance_embed') and \
|
|
self.model.guidance_embed is True:
|
|
guidance = torch.tensor([guidance_scale] * batch_size, device=device, dtype=dtype)
|
|
|
|
timesteps = self.scheduler.timesteps
|
|
|
|
for i, t in enumerate(timesteps):
|
|
# expand the latents if we are doing classifier free guidance
|
|
if do_classifier_free_guidance:
|
|
latent_model_input = torch.cat([latents] * 2)
|
|
else:
|
|
latent_model_input = latents
|
|
|
|
timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype)
|
|
timestep = timestep / self.scheduler.num_train_timesteps
|
|
noise_pred = self.model(latent_model_input, timestep, cond, guidance=guidance)
|
|
|
|
if do_classifier_free_guidance:
|
|
noise_pred_cond, noise_pred_uncond = noise_pred.chunk(2)
|
|
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
|
|
|
|
# compute the previous noisy sample x_t -> x_t-1
|
|
latents = self.scheduler.step(noise_pred, t, latents)
|
|
|
|
if callback is not None and i % callback_steps == 0:
|
|
step_idx = i // getattr(self.scheduler, "order", 1)
|
|
callback(step_idx, t, latents)
|
|
|
|
latents = 1. / self.vae.scale_factor * latents
|
|
mesh = self.vae.decode(latents, bounds = bounds, octree_res = octree_res, num_chunks = num_chunks)
|
|
|
|
return export_to_trimesh(mesh)
|
|
|