456 lines
14 KiB
Python
456 lines
14 KiB
Python
"""
|
|
Utilities for ComfyUI MV-Adapter nodes
|
|
"""
|
|
|
|
import os
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
from typing import Optional, Union, List
|
|
from torchvision import transforms
|
|
from transformers import AutoModelForImageSegmentation
|
|
|
|
from diffusers import AutoencoderKL, DDPMScheduler, LCMScheduler, UNet2DConditionModel
|
|
from huggingface_hub import snapshot_download
|
|
|
|
from .mvadapter.pipelines.pipeline_mvadapter_i2mv_sdxl import MVAdapterI2MVSDXLPipeline
|
|
from .mvadapter.pipelines.pipeline_mvadapter_t2mv_sdxl import MVAdapterT2MVSDXLPipeline
|
|
from .mvadapter.models.attention_processor import DecoupledMVRowColSelfAttnProcessor2_0
|
|
from .mvadapter.schedulers.scheduling_shift_snr import ShiftSNRScheduler
|
|
from .mvadapter.utils.mesh_utils import (
|
|
NVDiffRastContextWrapper,
|
|
get_orthogonal_camera,
|
|
load_mesh,
|
|
render,
|
|
)
|
|
from .mvadapter.utils import make_image_grid, tensor_to_image
|
|
from .mvadapter.utils.geometry import get_plucker_embeds_from_cameras_ortho
|
|
|
|
try:
|
|
from mmgp import offload, profile_type
|
|
MMGP_AVAILABLE = True
|
|
except ImportError:
|
|
MMGP_AVAILABLE = False
|
|
print("Warning: mmgp not available, memory optimization disabled")
|
|
|
|
|
|
def prepare_pipeline(
|
|
base_model: str,
|
|
vae_model: Optional[str] = None,
|
|
unet_model: Optional[str] = None,
|
|
lora_model: Optional[str] = None,
|
|
adapter_path: str = "huanngzh/mv-adapter",
|
|
scheduler: str = "default",
|
|
num_views: int = 6,
|
|
device: str = "cuda",
|
|
dtype: torch.dtype = torch.float16,
|
|
use_mmgp: bool = False,
|
|
adapter_local_path: Optional[str] = None,
|
|
):
|
|
"""
|
|
Prepare MV-Adapter pipeline
|
|
|
|
Args:
|
|
base_model: Base Stable Diffusion XL model
|
|
vae_model: Custom VAE model (optional)
|
|
unet_model: Custom UNet model (optional)
|
|
lora_model: LoRA model (optional)
|
|
adapter_path: Path to adapter weights
|
|
scheduler: Scheduler type ("default", "ddpm", "lcm")
|
|
num_views: Number of views to generate
|
|
device: Device for computation
|
|
dtype: Data type
|
|
use_mmgp: Use mmgp for memory optimization
|
|
adapter_local_path: Local path for adapter download
|
|
"""
|
|
# Load vae and unet if provided
|
|
pipe_kwargs = {}
|
|
if vae_model is not None:
|
|
pipe_kwargs["vae"] = AutoencoderKL.from_pretrained(vae_model)
|
|
if unet_model is not None:
|
|
pipe_kwargs["unet"] = UNet2DConditionModel.from_pretrained(unet_model)
|
|
|
|
# Prepare pipeline
|
|
pipe = MVAdapterI2MVSDXLPipeline.from_pretrained(base_model, **pipe_kwargs)
|
|
|
|
# Load scheduler if provided
|
|
if scheduler != "default":
|
|
scheduler_class = None
|
|
if scheduler == "ddpm":
|
|
scheduler_class = DDPMScheduler
|
|
elif scheduler == "lcm":
|
|
scheduler_class = LCMScheduler
|
|
|
|
pipe.scheduler = ShiftSNRScheduler.from_scheduler(
|
|
pipe.scheduler,
|
|
shift_mode="interpolated",
|
|
shift_scale=8.0,
|
|
scheduler_class=scheduler_class,
|
|
)
|
|
|
|
pipe.init_custom_adapter(
|
|
num_views=num_views, self_attn_processor=DecoupledMVRowColSelfAttnProcessor2_0
|
|
)
|
|
|
|
# Load adapter weights
|
|
weight_name = "mvadapter_ig2mv_sdxl.safetensors"
|
|
pipe.load_custom_adapter(adapter_path, weight_name=weight_name, local_cache_dir=adapter_local_path)
|
|
|
|
pipe.to(device=device, dtype=dtype)
|
|
pipe.cond_encoder.to(device=device, dtype=dtype)
|
|
|
|
# load lora if provided
|
|
if lora_model is not None:
|
|
model_, name_ = lora_model.rsplit("/", 1)
|
|
pipe.load_lora_weights(model_, weight_name=name_)
|
|
|
|
# Apply mmgp optimization if available and requested
|
|
if use_mmgp and MMGP_AVAILABLE:
|
|
import gc
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
all_models = {
|
|
"vae": pipe.vae,
|
|
"text_encoder": pipe.text_encoder,
|
|
"text_encoder_2": pipe.text_encoder_2,
|
|
"unet": pipe.unet,
|
|
"cond_encoder": pipe.cond_encoder,
|
|
}
|
|
|
|
if hasattr(pipe, 'image_encoder') and pipe.image_encoder is not None:
|
|
all_models["image_encoder"] = pipe.image_encoder
|
|
|
|
offload.profile(all_models, profile_type.HighRAM_LowVRAM)
|
|
print("mmgp profiling applied successfully.")
|
|
|
|
# Enable VAE slicing for memory efficiency
|
|
pipe.enable_vae_slicing()
|
|
|
|
return pipe
|
|
|
|
|
|
# def create_bg_remover(device: str = "cuda"):
|
|
# """Create background removal function"""
|
|
# birefnet = AutoModelForImageSegmentation.from_pretrained(
|
|
# "ZhengPeng7/BiRefNet", trust_remote_code=True
|
|
# )
|
|
# birefnet.to(device)
|
|
# transform_image = transforms.Compose([
|
|
# transforms.Resize((1024, 1024)),
|
|
# transforms.ToTensor(),
|
|
# transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
|
|
# ])
|
|
|
|
# def remove_bg(image):
|
|
# image_size = image.size
|
|
# input_images = transform_image(image).unsqueeze(0).to(device)
|
|
# with torch.no_grad():
|
|
# preds = birefnet(input_images)[-1].sigmoid().cpu()
|
|
# pred = preds[0].squeeze()
|
|
# pred_pil = transforms.ToPILImage()(pred)
|
|
# mask = pred_pil.resize(image_size)
|
|
# image.putalpha(mask)
|
|
# return image
|
|
|
|
# return remove_bg
|
|
|
|
|
|
def preprocess_image(image: Image.Image, height: int, width: int) -> Image.Image:
|
|
"""Preprocess image for MV-Adapter"""
|
|
image = np.array(image)
|
|
alpha = image[..., 3] > 0
|
|
H, W = alpha.shape
|
|
# get the bounding box of alpha
|
|
y, x = np.where(alpha)
|
|
y0, y1 = max(y.min() - 1, 0), min(y.max() + 1, H)
|
|
x0, x1 = max(x.min() - 1, 0), min(x.max() + 1, W)
|
|
image_center = image[y0:y1, x0:x1]
|
|
# resize the longer side to H * 0.9
|
|
H, W, _ = image_center.shape
|
|
if H > W:
|
|
W = int(W * (height * 0.9) / H)
|
|
H = int(height * 0.9)
|
|
else:
|
|
H = int(H * (width * 0.9) / W)
|
|
W = int(width * 0.9)
|
|
image_center = np.array(Image.fromarray(image_center).resize((W, H)))
|
|
# pad to H, W
|
|
start_h = (height - H) // 2
|
|
start_w = (width - W) // 2
|
|
image = np.zeros((height, width, 4), dtype=np.uint8)
|
|
image[start_h : start_h + H, start_w : start_w + W] = image_center
|
|
image = image.astype(np.float32) / 255.0
|
|
image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
|
|
image = (image * 255).clip(0, 255).astype(np.uint8)
|
|
image = Image.fromarray(image)
|
|
|
|
return image
|
|
|
|
|
|
def run_pipeline(
|
|
pipe,
|
|
mesh_path: str,
|
|
num_views: int,
|
|
text: str,
|
|
image: Union[str, Image.Image],
|
|
height: int,
|
|
width: int,
|
|
num_inference_steps: int,
|
|
guidance_scale: float,
|
|
seed: int,
|
|
remove_bg_fn=None,
|
|
reference_conditioning_scale: float = 1.0,
|
|
negative_prompt: str = "watermark, ugly, deformed, noisy, blurry, low contrast",
|
|
lora_scale: float = 1.0,
|
|
device: str = "cuda",
|
|
):
|
|
"""
|
|
Execute multi-view image generation
|
|
"""
|
|
# Prepare cameras
|
|
cameras = get_orthogonal_camera(
|
|
elevation_deg=[0, 0, 0, 0, 89.99, -89.99],
|
|
distance=[1.8] * num_views,
|
|
left=-0.55,
|
|
right=0.55,
|
|
bottom=-0.55,
|
|
top=0.55,
|
|
azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]],
|
|
device=device,
|
|
)
|
|
ctx = NVDiffRastContextWrapper(device=device)
|
|
|
|
mesh = load_mesh(mesh_path, rescale=True, device=device)
|
|
render_out = render(
|
|
ctx,
|
|
mesh,
|
|
cameras,
|
|
height=height,
|
|
width=width,
|
|
render_attr=False,
|
|
normal_background=0.0,
|
|
)
|
|
|
|
pos_images = tensor_to_image((render_out.pos + 0.5).clamp(0, 1), batched=True)
|
|
normal_images = tensor_to_image(
|
|
(render_out.normal / 2 + 0.5).clamp(0, 1), batched=True
|
|
)
|
|
|
|
control_images = (
|
|
torch.cat(
|
|
[
|
|
(render_out.pos + 0.5).clamp(0, 1),
|
|
(render_out.normal / 2 + 0.5).clamp(0, 1),
|
|
],
|
|
dim=-1,
|
|
)
|
|
.permute(0, 3, 1, 2)
|
|
.to(device)
|
|
)
|
|
|
|
# Prepare image
|
|
reference_image = Image.open(image) if isinstance(image, str) else image
|
|
# if remove_bg_fn is not None:
|
|
# reference_image = remove_bg_fn(reference_image)
|
|
# reference_image = preprocess_image(reference_image, height, width)
|
|
if reference_image.mode == "RGBA":
|
|
reference_image = preprocess_image(reference_image, height, width)
|
|
|
|
pipe_kwargs = {}
|
|
if seed != -1 and isinstance(seed, int):
|
|
pipe_kwargs["generator"] = torch.Generator(device=device).manual_seed(seed)
|
|
|
|
images = pipe(
|
|
text,
|
|
height=height,
|
|
width=width,
|
|
num_inference_steps=num_inference_steps,
|
|
guidance_scale=guidance_scale,
|
|
num_images_per_prompt=num_views,
|
|
control_image=control_images,
|
|
control_conditioning_scale=1.0,
|
|
reference_image=reference_image,
|
|
reference_conditioning_scale=reference_conditioning_scale,
|
|
negative_prompt=negative_prompt,
|
|
cross_attention_kwargs={"scale": lora_scale},
|
|
**pipe_kwargs,
|
|
).images
|
|
|
|
return images, pos_images, normal_images, reference_image
|
|
|
|
|
|
def create_image_grid(images: List[Image.Image], rows: int = 1) -> Image.Image:
|
|
"""Создание сетки изображений"""
|
|
return make_image_grid(images, rows=rows)
|
|
|
|
|
|
def pil_to_torch_tensor(image: Image.Image) -> torch.Tensor:
|
|
"""Конвертация PIL изображения в torch tensor"""
|
|
import numpy as np
|
|
array = np.array(image)
|
|
tensor = torch.from_numpy(array).float() / 255.0
|
|
|
|
# Add batch dimension if needed
|
|
if tensor.dim() == 3:
|
|
tensor = tensor.unsqueeze(0)
|
|
|
|
return tensor
|
|
|
|
|
|
def pils_to_torch_tensor(images: List[Image.Image]) -> torch.Tensor:
|
|
"""Конвертация списка PIL изображений в torch tensor"""
|
|
tensors = [pil_to_torch_tensor(img).squeeze(0) for img in images]
|
|
return torch.stack(tensors, dim=0)
|
|
|
|
|
|
def prepare_t2mv_pipeline(
|
|
base_model: str,
|
|
vae_model: Optional[str] = None,
|
|
unet_model: Optional[str] = None,
|
|
lora_model: Optional[str] = None,
|
|
adapter_path: str = "huanngzh/mv-adapter",
|
|
scheduler: str = "default",
|
|
num_views: int = 6,
|
|
device: str = "cuda",
|
|
dtype: torch.dtype = torch.float16,
|
|
use_mmgp: bool = False,
|
|
adapter_local_path: Optional[str] = None,
|
|
):
|
|
"""
|
|
Prepare MV-Adapter T2MV pipeline
|
|
|
|
Args:
|
|
base_model: Base Stable Diffusion XL model
|
|
vae_model: Custom VAE model (optional)
|
|
unet_model: Custom UNet model (optional)
|
|
lora_model: LoRA model (optional)
|
|
adapter_path: Path to adapter weights
|
|
scheduler: Scheduler type ("default", "ddpm", "lcm")
|
|
num_views: Number of views to generate
|
|
device: Device for computation
|
|
dtype: Data type
|
|
use_mmgp: Use mmgp for memory optimization
|
|
adapter_local_path: Local path for adapter download
|
|
"""
|
|
# Load vae and unet if provided
|
|
pipe_kwargs = {}
|
|
if vae_model is not None:
|
|
pipe_kwargs["vae"] = AutoencoderKL.from_pretrained(vae_model)
|
|
if unet_model is not None:
|
|
pipe_kwargs["unet"] = UNet2DConditionModel.from_pretrained(unet_model)
|
|
|
|
# Prepare pipeline
|
|
pipe = MVAdapterT2MVSDXLPipeline.from_pretrained(base_model, **pipe_kwargs)
|
|
|
|
# Load scheduler if provided
|
|
if scheduler != "default":
|
|
scheduler_class = None
|
|
if scheduler == "ddpm":
|
|
scheduler_class = DDPMScheduler
|
|
elif scheduler == "lcm":
|
|
scheduler_class = LCMScheduler
|
|
|
|
pipe.scheduler = ShiftSNRScheduler.from_scheduler(
|
|
pipe.scheduler,
|
|
shift_mode="interpolated",
|
|
shift_scale=8.0,
|
|
scheduler_class=scheduler_class,
|
|
)
|
|
|
|
pipe.init_custom_adapter(num_views=num_views)
|
|
|
|
# Load adapter weights
|
|
weight_name = "mvadapter_t2mv_sdxl.safetensors"
|
|
pipe.load_custom_adapter(adapter_path, weight_name=weight_name, local_cache_dir=adapter_local_path)
|
|
|
|
pipe.to(device=device, dtype=dtype)
|
|
pipe.cond_encoder.to(device=device, dtype=dtype)
|
|
|
|
# load lora if provided
|
|
if lora_model is not None:
|
|
model_, name_ = lora_model.rsplit("/", 1)
|
|
pipe.load_lora_weights(model_, weight_name=name_)
|
|
|
|
# Apply mmgp optimization if available and requested
|
|
if use_mmgp and MMGP_AVAILABLE:
|
|
import gc
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
all_models = {
|
|
"vae": pipe.vae,
|
|
"text_encoder": pipe.text_encoder,
|
|
"text_encoder_2": pipe.text_encoder_2,
|
|
"unet": pipe.unet,
|
|
"cond_encoder": pipe.cond_encoder,
|
|
}
|
|
|
|
if hasattr(pipe, 'image_encoder') and pipe.image_encoder is not None:
|
|
all_models["image_encoder"] = pipe.image_encoder
|
|
|
|
offload.profile(all_models, profile_type.HighRAM_LowVRAM)
|
|
print("mmgp profiling applied successfully.")
|
|
|
|
# Enable VAE slicing for memory efficiency
|
|
pipe.enable_vae_slicing()
|
|
|
|
return pipe
|
|
|
|
|
|
def run_t2mv_pipeline(
|
|
pipe,
|
|
num_views: int,
|
|
text: str,
|
|
height: int,
|
|
width: int,
|
|
num_inference_steps: int,
|
|
guidance_scale: float,
|
|
seed: int,
|
|
negative_prompt: str = "watermark, ugly, deformed, noisy, blurry, low contrast",
|
|
lora_scale: float = 1.0,
|
|
device: str = "cuda",
|
|
azimuth_deg: Optional[List[int]] = None,
|
|
):
|
|
"""
|
|
Execute text-to-multiview generation
|
|
"""
|
|
# Prepare cameras
|
|
if azimuth_deg is None:
|
|
azimuth_deg = [0, 45, 90, 180, 270, 315][:num_views]
|
|
|
|
cameras = get_orthogonal_camera(
|
|
elevation_deg=[0] * num_views,
|
|
distance=[1.8] * num_views,
|
|
left=-0.55,
|
|
right=0.55,
|
|
bottom=-0.55,
|
|
top=0.55,
|
|
azimuth_deg=[x - 90 for x in azimuth_deg],
|
|
device=device,
|
|
)
|
|
|
|
plucker_embeds = get_plucker_embeds_from_cameras_ortho(
|
|
cameras.c2w, [1.1] * num_views, width
|
|
)
|
|
control_images = ((plucker_embeds + 1.0) / 2.0).clamp(0, 1)
|
|
|
|
pipe_kwargs = {}
|
|
if seed != -1 and isinstance(seed, int):
|
|
pipe_kwargs["generator"] = torch.Generator(device=device).manual_seed(seed)
|
|
|
|
images = pipe(
|
|
text,
|
|
height=height,
|
|
width=width,
|
|
num_inference_steps=num_inference_steps,
|
|
guidance_scale=guidance_scale,
|
|
num_images_per_prompt=num_views,
|
|
control_image=control_images,
|
|
control_conditioning_scale=1.0,
|
|
negative_prompt=negative_prompt,
|
|
cross_attention_kwargs={"scale": lora_scale},
|
|
**pipe_kwargs,
|
|
).images
|
|
|
|
return images |