Spaces:
Running
on
L40S
Running
on
L40S
import spaces | |
from PIL import Image | |
import io | |
import argparse | |
import os | |
import random | |
import tempfile | |
from typing import Dict, Optional, Tuple | |
from omegaconf import OmegaConf | |
import numpy as np | |
import torch | |
from diffusers import AutoencoderKL, DDIMScheduler | |
from diffusers.utils import check_min_version | |
from tqdm.auto import tqdm | |
from transformers import CLIPTextModel, CLIPTokenizer, CLIPImageProcessor, CLIPVisionModelWithProjection | |
from torchvision import transforms | |
from canonicalize.models.unet_mv2d_condition import UNetMV2DConditionModel | |
from canonicalize.models.unet_mv2d_ref import UNetMV2DRefModel | |
from canonicalize.pipeline_canonicalize import CanonicalizationPipeline | |
from einops import rearrange | |
from torchvision.utils import save_image | |
import json | |
import cv2 | |
import onnxruntime as rt | |
from huggingface_hub.file_download import hf_hub_download | |
from huggingface_hub import list_repo_files | |
from rm_anime_bg.cli import get_mask, SCALE | |
import argparse | |
import os | |
import cv2 | |
import glob | |
import numpy as np | |
import matplotlib.pyplot as plt | |
from typing import Dict, Optional, List | |
from omegaconf import OmegaConf, DictConfig | |
from PIL import Image | |
from pathlib import Path | |
from dataclasses import dataclass | |
from typing import Dict | |
import torch | |
import torch.nn.functional as F | |
import torch.utils.checkpoint | |
import torchvision.transforms.functional as TF | |
from torch.utils.data import Dataset, DataLoader | |
from torchvision import transforms | |
from torchvision.utils import make_grid, save_image | |
from accelerate.utils import set_seed | |
from tqdm.auto import tqdm | |
from einops import rearrange, repeat | |
from multiview.pipeline_multiclass import StableUnCLIPImg2ImgPipeline | |
import os | |
import imageio | |
import numpy as np | |
import torch | |
import cv2 | |
import glob | |
import matplotlib.pyplot as plt | |
from PIL import Image | |
from torchvision.transforms import v2 | |
from pytorch_lightning import seed_everything | |
from omegaconf import OmegaConf | |
from tqdm import tqdm | |
from slrm.utils.train_util import instantiate_from_config | |
from slrm.utils.camera_util import ( | |
FOV_to_intrinsics, | |
get_circular_camera_poses, | |
) | |
from slrm.utils.mesh_util import save_obj, save_glb | |
from slrm.utils.infer_util import images_to_video | |
import cv2 | |
import numpy as np | |
import os | |
import trimesh | |
import argparse | |
import torch | |
import scipy | |
from PIL import Image | |
from refine.mesh_refine import geo_refine | |
from refine.func import make_star_cameras_orthographic | |
from refine.render import NormalsRenderer, calc_vertex_normals | |
import pytorch3d | |
from pytorch3d.structures import Meshes | |
from sklearn.neighbors import KDTree | |
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry | |
check_min_version("0.24.0") | |
weight_dtype = torch.float16 | |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
VIEWS = ['front', 'front_right', 'right', 'back', 'left', 'front_left'] | |
def set_seed(seed): | |
random.seed(seed) | |
np.random.seed(seed) | |
torch.manual_seed(seed) | |
torch.cuda.manual_seed_all(seed) | |
session_infer_path = hf_hub_download( | |
repo_id="skytnt/anime-seg", filename="isnetis.onnx", | |
) | |
providers: list[str] = ["CPUExecutionProvider"] | |
if "CUDAExecutionProvider" in rt.get_available_providers(): | |
providers = ["CUDAExecutionProvider"] | |
bkg_remover_session_infer = rt.InferenceSession( | |
session_infer_path, providers=providers, | |
) | |
def remove_background( | |
img: np.ndarray, | |
alpha_min: float, | |
alpha_max: float, | |
) -> list: | |
img = np.array(img) | |
mask = get_mask(bkg_remover_session_infer, img) | |
mask[mask < alpha_min] = 0.0 | |
mask[mask > alpha_max] = 1.0 | |
img_after = (mask * img).astype(np.uint8) | |
mask = (mask * SCALE).astype(np.uint8) | |
img_after = np.concatenate([img_after, mask], axis=2, dtype=np.uint8) | |
return Image.fromarray(img_after) | |
def process_image(image, totensor, width, height): | |
assert image.mode == "RGBA" | |
# Find non-transparent pixels | |
non_transparent = np.nonzero(np.array(image)[..., 3]) | |
min_x, max_x = non_transparent[1].min(), non_transparent[1].max() | |
min_y, max_y = non_transparent[0].min(), non_transparent[0].max() | |
image = image.crop((min_x, min_y, max_x, max_y)) | |
# paste to center | |
max_dim = max(image.width, image.height) | |
max_height = int(max_dim * 1.2) | |
max_width = int(max_dim / (height/width) * 1.2) | |
new_image = Image.new("RGBA", (max_width, max_height)) | |
left = (max_width - image.width) // 2 | |
top = (max_height - image.height) // 2 | |
new_image.paste(image, (left, top)) | |
image = new_image.resize((width, height), resample=Image.BICUBIC) | |
image = np.array(image) | |
image = image.astype(np.float32) / 255. | |
assert image.shape[-1] == 4 # RGBA | |
alpha = image[..., 3:4] | |
bg_color = np.array([1., 1., 1.], dtype=np.float32) | |
image = image[..., :3] * alpha + bg_color * (1 - alpha) | |
return totensor(image) | |
def inference(validation_pipeline, input_image, vae, feature_extractor, image_encoder, unet, ref_unet, tokenizer, | |
text_encoder, pretrained_model_path, validation, val_width, val_height, unet_condition_type, | |
use_noise=True, noise_d=256, crop=False, seed=100, timestep=20): | |
set_seed(seed) | |
generator = torch.Generator(device=device).manual_seed(seed) | |
totensor = transforms.ToTensor() | |
prompts = "high quality, best quality" | |
prompt_ids = tokenizer( | |
prompts, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, | |
return_tensors="pt" | |
).input_ids[0] | |
# (B*Nv, 3, H, W) | |
B = 1 | |
if input_image.mode != "RGBA": | |
# remove background | |
input_image = remove_background(input_image, 0.1, 0.9) | |
imgs_in = process_image(input_image, totensor, val_width, val_height) | |
imgs_in = rearrange(imgs_in.unsqueeze(0).unsqueeze(0), "B Nv C H W -> (B Nv) C H W") | |
with torch.autocast('cuda' if torch.cuda.is_available() else 'cpu', dtype=weight_dtype): | |
imgs_in = imgs_in.to(device=device) | |
# B*Nv images | |
out = validation_pipeline(prompt=prompts, image=imgs_in.to(weight_dtype), generator=generator, | |
num_inference_steps=timestep, prompt_ids=prompt_ids, | |
height=val_height, width=val_width, unet_condition_type=unet_condition_type, | |
use_noise=use_noise, **validation,) | |
out = rearrange(out, "B C f H W -> (B f) C H W", f=1) | |
print("OUT!!!!!!") | |
img_buf = io.BytesIO() | |
save_image(out[0], img_buf, format='PNG') | |
img_buf.seek(0) | |
img = Image.open(img_buf) | |
print("OUT2!!!!!!") | |
torch.cuda.empty_cache() | |
return img | |
######### Multi View Part ############# | |
weight_dtype = torch.float16 | |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
def tensor_to_numpy(tensor): | |
return tensor.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy() | |
class TestConfig: | |
pretrained_model_name_or_path: str | |
pretrained_unet_path:Optional[str] | |
revision: Optional[str] | |
validation_dataset: Dict | |
save_dir: str | |
seed: Optional[int] | |
validation_batch_size: int | |
dataloader_num_workers: int | |
save_mode: str | |
local_rank: int | |
pipe_kwargs: Dict | |
pipe_validation_kwargs: Dict | |
unet_from_pretrained_kwargs: Dict | |
validation_grid_nrow: int | |
camera_embedding_lr_mult: float | |
num_views: int | |
camera_embedding_type: str | |
pred_type: str | |
regress_elevation: bool | |
enable_xformers_memory_efficient_attention: bool | |
cond_on_normals: bool | |
cond_on_colors: bool | |
regress_elevation: bool | |
regress_focal_length: bool | |
def convert_to_numpy(tensor): | |
return tensor.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy() | |
def save_image(tensor): | |
ndarr = convert_to_numpy(tensor) | |
return save_image_numpy(ndarr) | |
def save_image_numpy(ndarr): | |
im = Image.fromarray(ndarr) | |
# pad to square | |
if im.size[0] != im.size[1]: | |
size = max(im.size) | |
new_im = Image.new("RGB", (size, size)) | |
# set to white | |
new_im.paste((255, 255, 255), (0, 0, size, size)) | |
new_im.paste(im, ((size - im.size[0]) // 2, (size - im.size[1]) // 2)) | |
im = new_im | |
# resize to 1024x1024 | |
im = im.resize((1024, 1024), Image.LANCZOS) | |
return im | |
def run_multiview_infer(data, pipeline, cfg: TestConfig, num_levels=3): | |
if cfg.seed is None: | |
generator = None | |
else: | |
generator = torch.Generator(device=pipeline.unet.device).manual_seed(cfg.seed) | |
images_cond = [] | |
results = {} | |
torch.cuda.empty_cache() | |
images_cond.append(data['image_cond_rgb'][:, 0].cuda()) | |
imgs_in = torch.cat([data['image_cond_rgb']]*2, dim=0).cuda() | |
num_views = imgs_in.shape[1] | |
imgs_in = rearrange(imgs_in, "B Nv C H W -> (B Nv) C H W")# (B*Nv, 3, H, W) | |
target_h, target_w = imgs_in.shape[-2], imgs_in.shape[-1] | |
normal_prompt_embeddings, clr_prompt_embeddings = data['normal_prompt_embeddings'].cuda(), data['color_prompt_embeddings'].cuda() | |
prompt_embeddings = torch.cat([normal_prompt_embeddings, clr_prompt_embeddings], dim=0) | |
prompt_embeddings = rearrange(prompt_embeddings, "B Nv N C -> (B Nv) N C") | |
# B*Nv images | |
unet_out = pipeline( | |
imgs_in, None, prompt_embeds=prompt_embeddings, | |
generator=generator, guidance_scale=3.0, output_type='pt', num_images_per_prompt=1, | |
height=cfg.height, width=cfg.width, | |
num_inference_steps=40, eta=1.0, | |
num_levels=num_levels, | |
) | |
for level in range(num_levels): | |
out = unet_out[level].images | |
bsz = out.shape[0] // 2 | |
normals_pred = out[:bsz] | |
images_pred = out[bsz:] | |
if num_levels == 2: | |
results[level+1] = {'normals': [], 'images': []} | |
else: | |
results[level] = {'normals': [], 'images': []} | |
for i in range(bsz//num_views): | |
img_in_ = images_cond[-1][i].to(out.device) | |
for j in range(num_views): | |
view = VIEWS[j] | |
idx = i*num_views + j | |
normal = normals_pred[idx] | |
color = images_pred[idx] | |
## save color and normal--------------------- | |
new_normal = save_image(normal) | |
new_color = save_image(color) | |
if num_levels == 2: | |
results[level+1]['normals'].append(new_normal) | |
results[level+1]['images'].append(new_color) | |
else: | |
results[level]['normals'].append(new_normal) | |
results[level]['images'].append(new_color) | |
torch.cuda.empty_cache() | |
return results | |
def load_multiview_pipeline(cfg): | |
pipeline = StableUnCLIPImg2ImgPipeline.from_pretrained( | |
cfg.pretrained_path, | |
torch_dtype=torch.float16,) | |
pipeline.unet.enable_xformers_memory_efficient_attention() | |
if torch.cuda.is_available(): | |
pipeline.to(device) | |
return pipeline | |
class InferAPI: | |
def __init__(self, | |
canonical_configs, | |
multiview_configs, | |
slrm_configs, | |
refine_configs): | |
self.canonical_configs = canonical_configs | |
self.multiview_configs = multiview_configs | |
self.slrm_configs = slrm_configs | |
self.refine_configs = refine_configs | |
repo_id = "hyz317/StdGEN" | |
all_files = list_repo_files(repo_id, revision="main") | |
for file in all_files: | |
if os.path.exists(file): | |
continue | |
hf_hub_download(repo_id, file, local_dir="./ckpt") | |
self.canonical_infer = InferCanonicalAPI(self.canonical_configs) | |
# self.multiview_infer = InferMultiviewAPI(self.multiview_configs) | |
# self.slrm_infer = InferSlrmAPI(self.slrm_configs) | |
# self.refine_infer = InferRefineAPI(self.refine_configs) | |
def genStage1(self, img, seed): | |
return self.canonical_infer.gen(img, seed) | |
def genStage2(self, img, seed, num_levels): | |
return self.multiview_infer.gen(img, seed, num_levels) | |
def genStage3(self, img): | |
return self.slrm_infer.gen(img) | |
def genStage4(self, meshes, imgs): | |
return self.refine_infer.refine(meshes, imgs) | |
############## Refine ############## | |
def fix_vert_color_glb(mesh_path): | |
from pygltflib import GLTF2, Material, PbrMetallicRoughness | |
obj1 = GLTF2().load(mesh_path) | |
obj1.meshes[0].primitives[0].material = 0 | |
obj1.materials.append(Material( | |
pbrMetallicRoughness = PbrMetallicRoughness( | |
baseColorFactor = [1.0, 1.0, 1.0, 1.0], | |
metallicFactor = 0., | |
roughnessFactor = 1.0, | |
), | |
emissiveFactor = [0.0, 0.0, 0.0], | |
doubleSided = True, | |
)) | |
obj1.save(mesh_path) | |
def srgb_to_linear(c_srgb): | |
c_linear = np.where(c_srgb <= 0.04045, c_srgb / 12.92, ((c_srgb + 0.055) / 1.055) ** 2.4) | |
return c_linear.clip(0, 1.) | |
def save_py3dmesh_with_trimesh_fast(meshes: Meshes, save_glb_path, apply_sRGB_to_LinearRGB=True): | |
# convert from pytorch3d meshes to trimesh mesh | |
vertices = meshes.verts_packed().cpu().float().numpy() | |
triangles = meshes.faces_packed().cpu().long().numpy() | |
np_color = meshes.textures.verts_features_packed().cpu().float().numpy() | |
if save_glb_path.endswith(".glb"): | |
# rotate 180 along +Y | |
vertices[:, [0, 2]] = -vertices[:, [0, 2]] | |
if apply_sRGB_to_LinearRGB: | |
np_color = srgb_to_linear(np_color) | |
assert vertices.shape[0] == np_color.shape[0] | |
assert np_color.shape[1] == 3 | |
assert 0 <= np_color.min() and np_color.max() <= 1.001, f"min={np_color.min()}, max={np_color.max()}" | |
np_color = np.clip(np_color, 0, 1) | |
mesh = trimesh.Trimesh(vertices=vertices, faces=triangles, vertex_colors=np_color) | |
mesh.remove_unreferenced_vertices() | |
# save mesh | |
mesh.export(save_glb_path) | |
if save_glb_path.endswith(".glb"): | |
fix_vert_color_glb(save_glb_path) | |
print(f"saving to {save_glb_path}") | |
def calc_horizontal_offset(target_img, source_img): | |
target_mask = target_img.astype(np.float32).sum(axis=-1) > 750 | |
source_mask = source_img.astype(np.float32).sum(axis=-1) > 750 | |
best_offset = -114514 | |
for offset in range(-200, 200): | |
offset_mask = np.roll(source_mask, offset, axis=1) | |
overlap = (target_mask & offset_mask).sum() | |
if overlap > best_offset: | |
best_offset = overlap | |
best_offset_value = offset | |
return best_offset_value | |
def calc_horizontal_offset2(target_mask, source_img): | |
source_mask = source_img.astype(np.float32).sum(axis=-1) > 750 | |
best_offset = -114514 | |
for offset in range(-200, 200): | |
offset_mask = np.roll(source_mask, offset, axis=1) | |
overlap = (target_mask & offset_mask).sum() | |
if overlap > best_offset: | |
best_offset = overlap | |
best_offset_value = offset | |
return best_offset_value | |
def get_distract_mask(generator, color_0, color_1, normal_0=None, normal_1=None, thres=0.25, ratio=0.50, outside_thres=0.10, outside_ratio=0.20): | |
distract_area = np.abs(color_0 - color_1).sum(axis=-1) > thres | |
if normal_0 is not None and normal_1 is not None: | |
distract_area |= np.abs(normal_0 - normal_1).sum(axis=-1) > thres | |
labeled_array, num_features = scipy.ndimage.label(distract_area) | |
results = [] | |
random_sampled_points = [] | |
for i in range(num_features + 1): | |
if np.sum(labeled_array == i) > 1000 and np.sum(labeled_array == i) < 100000: | |
results.append((i, np.sum(labeled_array == i))) | |
# random sample a point in the area | |
points = np.argwhere(labeled_array == i) | |
random_sampled_points.append(points[np.random.randint(0, points.shape[0])]) | |
results = sorted(results, key=lambda x: x[1], reverse=True) # [1:] | |
distract_mask = np.zeros_like(distract_area) | |
distract_bbox = np.zeros_like(distract_area) | |
for i, _ in results: | |
distract_mask |= labeled_array == i | |
bbox = np.argwhere(labeled_array == i) | |
min_x, min_y = bbox.min(axis=0) | |
max_x, max_y = bbox.max(axis=0) | |
distract_bbox[min_x:max_x, min_y:max_y] = 1 | |
points = np.array(random_sampled_points)[:, ::-1] | |
labels = np.ones(len(points), dtype=np.int32) | |
masks = generator.generate((color_1 * 255).astype(np.uint8)) | |
outside_area = np.abs(color_0 - color_1).sum(axis=-1) < outside_thres | |
final_mask = np.zeros_like(distract_mask) | |
for iii, mask in enumerate(masks): | |
mask['segmentation'] = cv2.resize(mask['segmentation'].astype(np.float32), (1024, 1024)) > 0.5 | |
intersection = np.logical_and(mask['segmentation'], distract_mask).sum() | |
total = mask['segmentation'].sum() | |
iou = intersection / total | |
outside_intersection = np.logical_and(mask['segmentation'], outside_area).sum() | |
outside_total = mask['segmentation'].sum() | |
outside_iou = outside_intersection / outside_total | |
if iou > ratio and outside_iou < outside_ratio: | |
final_mask |= mask['segmentation'] | |
# calculate coverage | |
intersection = np.logical_and(final_mask, distract_mask).sum() | |
total = distract_mask.sum() | |
coverage = intersection / total | |
if coverage < 0.8: | |
# use original distract mask | |
final_mask = (distract_mask.copy() * 255).astype(np.uint8) | |
final_mask = cv2.dilate(final_mask, np.ones((3, 3), np.uint8), iterations=3) | |
labeled_array_dilate, num_features_dilate = scipy.ndimage.label(final_mask) | |
for i in range(num_features_dilate + 1): | |
if np.sum(labeled_array_dilate == i) < 200: | |
final_mask[labeled_array_dilate == i] = 255 | |
final_mask = cv2.erode(final_mask, np.ones((3, 3), np.uint8), iterations=3) | |
final_mask = final_mask > 127 | |
return distract_mask, distract_bbox, random_sampled_points, final_mask | |
class InferRefineAPI: | |
def __init__(self, config): | |
self.sam = sam_model_registry["vit_h"](checkpoint="./ckpt/sam_vit_h_4b8939.pth").cuda() | |
self.generator = SamAutomaticMaskGenerator( | |
model=self.sam, | |
points_per_side=64, | |
pred_iou_thresh=0.80, | |
stability_score_thresh=0.92, | |
crop_n_layers=1, | |
crop_n_points_downscale_factor=2, | |
min_mask_region_area=100, | |
) | |
self.outside_ratio = 0.20 | |
def refine(self, meshes, imgs): | |
fixed_v, fixed_f, fixed_t = None, None, None | |
flow_vert, flow_vector = None, None | |
last_colors, last_normals = None, None | |
last_front_color, last_front_normal = None, None | |
distract_mask = None | |
mv, proj = make_star_cameras_orthographic(8, 1, r=1.2) | |
mv = mv[[4, 3, 2, 0, 6, 5]] | |
renderer = NormalsRenderer(mv,proj,(1024,1024)) | |
results = [] | |
for name_idx, level in zip([2, 0, 1], [2, 1, 0]): | |
mesh = trimesh.load(meshes[name_idx]) | |
new_mesh = mesh.split(only_watertight=False) | |
new_mesh = [ j for j in new_mesh if len(j.vertices) >= 300 ] | |
mesh = trimesh.Scene(new_mesh).dump(concatenate=True) | |
mesh_v, mesh_f = mesh.vertices, mesh.faces | |
if last_colors is None: | |
images = renderer.render( | |
torch.tensor(mesh_v, device='cuda').float(), | |
torch.ones_like(torch.from_numpy(mesh_v), device='cuda').float(), | |
torch.tensor(mesh_f, device='cuda'), | |
) | |
mask = (images[..., 3] < 0.9).cpu().numpy() | |
colors, normals = [], [] | |
for i in range(6): | |
color = np.array(imgs[level]['images'][i]) | |
normal = np.array(imgs[level]['normals'][i]) | |
if last_colors is not None: | |
offset = calc_horizontal_offset(np.array(last_colors[i]), color) | |
# print('offset', i, offset) | |
else: | |
offset = calc_horizontal_offset2(mask[i], color) | |
# print('init offset', i, offset) | |
if offset != 0: | |
color = np.roll(color, offset, axis=1) | |
normal = np.roll(normal, offset, axis=1) | |
color = Image.fromarray(color) | |
normal = Image.fromarray(normal) | |
colors.append(color) | |
normals.append(normal) | |
if last_front_color is not None and level == 0: | |
original_mask, distract_bbox, _, distract_mask = get_distract_mask(self.generator, last_front_color, np.array(colors[0]).astype(np.float32) / 255.0, outside_ratio=self.outside_ratio) | |
else: | |
distract_mask = None | |
distract_bbox = None | |
last_front_color = np.array(colors[0]).astype(np.float32) / 255.0 | |
last_front_normal = np.array(normals[0]).astype(np.float32) / 255.0 | |
if last_colors is None: | |
from copy import deepcopy | |
last_colors, last_normals = deepcopy(colors), deepcopy(normals) | |
# my mesh flow weight by nearest vertexs | |
if fixed_v is not None and fixed_f is not None and level == 1: | |
t = trimesh.Trimesh(vertices=mesh_v, faces=mesh_f) | |
fixed_v_cpu = fixed_v.cpu().numpy() | |
kdtree_anchor = KDTree(fixed_v_cpu) | |
kdtree_mesh_v = KDTree(mesh_v) | |
_, idx_anchor = kdtree_anchor.query(mesh_v, k=1) | |
_, idx_mesh_v = kdtree_mesh_v.query(mesh_v, k=25) | |
idx_anchor = idx_anchor.squeeze() | |
neighbors = torch.tensor(mesh_v).cuda()[idx_mesh_v] # V, 25, 3 | |
# calculate the distances neighbors [V, 25, 3]; mesh_v [V, 3] -> [V, 25] | |
neighbor_dists = torch.norm(neighbors - torch.tensor(mesh_v).cuda()[:, None], dim=-1) | |
neighbor_dists[neighbor_dists > 0.06] = 114514. | |
neighbor_weights = torch.exp(-neighbor_dists * 1.) | |
neighbor_weights = neighbor_weights / neighbor_weights.sum(dim=1, keepdim=True) | |
anchors = fixed_v[idx_anchor] # V, 3 | |
anchor_normals = calc_vertex_normals(fixed_v, fixed_f)[idx_anchor] # V, 3 | |
dis_anchor = torch.clamp(((anchors - torch.tensor(mesh_v).cuda()) * anchor_normals).sum(-1), min=0) + 0.01 | |
vec_anchor = dis_anchor[:, None] * anchor_normals # V, 3 | |
vec_anchor = vec_anchor[idx_mesh_v] # V, 25, 3 | |
weighted_vec_anchor = (vec_anchor * neighbor_weights[:, :, None]).sum(1) # V, 3 | |
mesh_v += weighted_vec_anchor.cpu().numpy() | |
t = trimesh.Trimesh(vertices=mesh_v, faces=mesh_f) | |
mesh_v = torch.tensor(mesh_v, device='cuda', dtype=torch.float32) | |
mesh_f = torch.tensor(mesh_f, device='cuda') | |
new_mesh, simp_v, simp_f = geo_refine(mesh_v, mesh_f, colors, normals, fixed_v=fixed_v, fixed_f=fixed_f, distract_mask=distract_mask, distract_bbox=distract_bbox) | |
# my mesh flow weight by nearest vertexs | |
try: | |
if fixed_v is not None and fixed_f is not None and level != 0: | |
new_mesh_v = new_mesh.verts_packed().cpu().numpy() | |
fixed_v_cpu = fixed_v.cpu().numpy() | |
kdtree_anchor = KDTree(fixed_v_cpu) | |
kdtree_mesh_v = KDTree(new_mesh_v) | |
_, idx_anchor = kdtree_anchor.query(new_mesh_v, k=1) | |
_, idx_mesh_v = kdtree_mesh_v.query(new_mesh_v, k=25) | |
idx_anchor = idx_anchor.squeeze() | |
neighbors = torch.tensor(new_mesh_v).cuda()[idx_mesh_v] # V, 25, 3 | |
# calculate the distances neighbors [V, 25, 3]; new_mesh_v [V, 3] -> [V, 25] | |
neighbor_dists = torch.norm(neighbors - torch.tensor(new_mesh_v).cuda()[:, None], dim=-1) | |
neighbor_dists[neighbor_dists > 0.06] = 114514. | |
neighbor_weights = torch.exp(-neighbor_dists * 1.) | |
neighbor_weights = neighbor_weights / neighbor_weights.sum(dim=1, keepdim=True) | |
anchors = fixed_v[idx_anchor] # V, 3 | |
anchor_normals = calc_vertex_normals(fixed_v, fixed_f)[idx_anchor] # V, 3 | |
dis_anchor = torch.clamp(((anchors - torch.tensor(new_mesh_v).cuda()) * anchor_normals).sum(-1), min=0) + 0.01 | |
vec_anchor = dis_anchor[:, None] * anchor_normals # V, 3 | |
vec_anchor = vec_anchor[idx_mesh_v] # V, 25, 3 | |
weighted_vec_anchor = (vec_anchor * neighbor_weights[:, :, None]).sum(1) # V, 3 | |
new_mesh_v += weighted_vec_anchor.cpu().numpy() | |
# replace new_mesh verts with new_mesh_v | |
new_mesh = Meshes(verts=[torch.tensor(new_mesh_v, device='cuda')], faces=new_mesh.faces_list(), textures=new_mesh.textures) | |
except Exception as e: | |
pass | |
notsimp_v, notsimp_f, notsimp_t = new_mesh.verts_packed(), new_mesh.faces_packed(), new_mesh.textures.verts_features_packed() | |
if fixed_v is None: | |
fixed_v, fixed_f = simp_v, simp_f | |
complete_v, complete_f, complete_t = notsimp_v, notsimp_f, notsimp_t | |
else: | |
fixed_f = torch.cat([fixed_f, simp_f + fixed_v.shape[0]], dim=0) | |
fixed_v = torch.cat([fixed_v, simp_v], dim=0) | |
complete_f = torch.cat([complete_f, notsimp_f + complete_v.shape[0]], dim=0) | |
complete_v = torch.cat([complete_v, notsimp_v], dim=0) | |
complete_t = torch.cat([complete_t, notsimp_t], dim=0) | |
if level == 2: | |
new_mesh = Meshes(verts=[new_mesh.verts_packed()], faces=[new_mesh.faces_packed()], textures=pytorch3d.renderer.mesh.textures.TexturesVertex(verts_features=[torch.ones_like(new_mesh.textures.verts_features_packed(), device=new_mesh.verts_packed().device)*0.5])) | |
save_py3dmesh_with_trimesh_fast(new_mesh, meshes[name_idx].replace('.obj', '_refined.obj'), apply_sRGB_to_LinearRGB=False) | |
results.append(meshes[name_idx].replace('.obj', '_refined.obj')) | |
# save whole mesh | |
save_py3dmesh_with_trimesh_fast(Meshes(verts=[complete_v], faces=[complete_f], textures=pytorch3d.renderer.mesh.textures.TexturesVertex(verts_features=[complete_t])), meshes[name_idx].replace('.obj', '_refined_whole.obj'), apply_sRGB_to_LinearRGB=False) | |
results.append(meshes[name_idx].replace('.obj', '_refined_whole.obj')) | |
return results | |
class InferSlrmAPI: | |
def __init__(self, config): | |
self.config_path = config['config_path'] | |
self.config = OmegaConf.load(self.config_path) | |
self.config_name = os.path.basename(self.config_path).replace('.yaml', '') | |
self.model_config = self.config.model_config | |
self.infer_config = self.config.infer_config | |
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
self.model = instantiate_from_config(self.model_config) | |
state_dict = torch.load(self.infer_config.model_path, map_location='cpu') | |
self.model.load_state_dict(state_dict, strict=False) | |
self.model = self.model.to(self.device) | |
self.model.init_flexicubes_geometry(self.device, fovy=30.0, is_ortho=self.model.is_ortho) | |
self.model = self.model.eval() | |
def gen(self, imgs): | |
imgs = [ cv2.imread(img[0])[:, :, ::-1] for img in imgs ] | |
imgs = np.stack(imgs, axis=0).astype(np.float32) / 255.0 | |
imgs = torch.from_numpy(np.array(imgs)).permute(0, 3, 1, 2).contiguous().float() # (6, 3, 1024, 1024) | |
mesh_glb_fpaths = self.make3d(imgs) | |
return mesh_glb_fpaths[1:4] + mesh_glb_fpaths[0:1] | |
def make3d(self, images): | |
input_cameras = torch.tensor(np.load('slrm/cameras.npy')).to(device) | |
images = images.unsqueeze(0).to(device) | |
images = v2.functional.resize(images, (320, 320), interpolation=3, antialias=True).clamp(0, 1) | |
mesh_fpath = tempfile.NamedTemporaryFile(suffix=f".obj", delete=False).name | |
print(mesh_fpath) | |
mesh_basename = os.path.basename(mesh_fpath).split('.')[0] | |
mesh_dirname = os.path.dirname(mesh_fpath) | |
with torch.no_grad(): | |
# get triplane | |
planes = self.model.forward_planes(images, input_cameras.float()) | |
# get mesh | |
mesh_glb_fpaths = [] | |
for j in range(4): | |
mesh_glb_fpath = self.make_mesh(mesh_fpath.replace(mesh_fpath[-4:], f'_{j}{mesh_fpath[-4:]}'), planes, level=[0, 3, 4, 2][j]) | |
mesh_glb_fpaths.append(mesh_glb_fpath) | |
return mesh_glb_fpaths | |
def make_mesh(self, mesh_fpath, planes, level=None): | |
mesh_basename = os.path.basename(mesh_fpath).split('.')[0] | |
mesh_dirname = os.path.dirname(mesh_fpath) | |
mesh_glb_fpath = os.path.join(mesh_dirname, f"{mesh_basename}.glb") | |
with torch.no_grad(): | |
# get mesh | |
mesh_out = self.model.extract_mesh( | |
planes, | |
use_texture_map=False, | |
levels=torch.tensor([level]).to(device), | |
**self.infer_config, | |
) | |
vertices, faces, vertex_colors = mesh_out | |
vertices = vertices[:, [1, 2, 0]] | |
if level == 2: | |
# fill all vertex_colors with 127 | |
vertex_colors = np.ones_like(vertex_colors) * 127 | |
save_obj(vertices, faces, vertex_colors, mesh_fpath) | |
return mesh_fpath | |
class InferMultiviewAPI: | |
def __init__(self, config): | |
parser = argparse.ArgumentParser() | |
parser.add_argument("--seed", type=int, default=42) | |
parser.add_argument("--num_views", type=int, default=6) | |
parser.add_argument("--num_levels", type=int, default=3) | |
parser.add_argument("--pretrained_path", type=str, default='./ckpt/StdGEN-multiview-1024') | |
parser.add_argument("--height", type=int, default=1024) | |
parser.add_argument("--width", type=int, default=576) | |
self.cfg = parser.parse_args() | |
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
self.pipeline = load_multiview_pipeline(self.cfg) | |
self.results = {} | |
if torch.cuda.is_available(): | |
self.pipeline.to(device) | |
self.image_transforms = [transforms.Resize(int(max(self.cfg.height, self.cfg.width))), | |
transforms.CenterCrop((self.cfg.height, self.cfg.width)), | |
transforms.ToTensor(), | |
transforms.Lambda(lambda x: x * 2. - 1), | |
] | |
self.image_transforms = transforms.Compose(self.image_transforms) | |
prompt_embeds_path = './multiview/fixed_prompt_embeds_6view' | |
self.normal_text_embeds = torch.load(f'{prompt_embeds_path}/normal_embeds.pt') | |
self.color_text_embeds = torch.load(f'{prompt_embeds_path}/clr_embeds.pt') | |
self.total_views = self.cfg.num_views | |
def process_im(self, im): | |
im = self.image_transforms(im) | |
return im | |
def gen(self, img, seed, num_levels): | |
set_seed(seed) | |
data = {} | |
cond_im_rgb = self.process_im(img) | |
cond_im_rgb = torch.stack([cond_im_rgb] * self.total_views, dim=0) | |
data["image_cond_rgb"] = cond_im_rgb[None, ...] | |
data["normal_prompt_embeddings"] = self.normal_text_embeds[None, ...] | |
data["color_prompt_embeddings"] = self.color_text_embeds[None, ...] | |
results = run_multiview_infer(data, self.pipeline, self.cfg, num_levels=num_levels) | |
for k in results: | |
self.results[k] = results[k] | |
return results | |
class InferCanonicalAPI: | |
def __init__(self, config): | |
self.config = config | |
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
self.config_path = config['config_path'] | |
self.loaded_config = OmegaConf.load(self.config_path) | |
self.setup(**self.loaded_config) | |
def setup(self, | |
validation: Dict, | |
pretrained_model_path: str, | |
local_crossattn: bool = True, | |
unet_from_pretrained_kwargs=None, | |
unet_condition_type=None, | |
use_noise=True, | |
noise_d=256, | |
timestep: int = 40, | |
width_input: int = 640, | |
height_input: int = 1024, | |
): | |
self.width_input = width_input | |
self.height_input = height_input | |
self.timestep = timestep | |
self.use_noise = use_noise | |
self.noise_d = noise_d | |
self.validation = validation | |
self.unet_condition_type = unet_condition_type | |
self.pretrained_model_path = pretrained_model_path | |
self.tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer") | |
self.text_encoder = CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder") | |
self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(pretrained_model_path, subfolder="image_encoder") | |
self.feature_extractor = CLIPImageProcessor() | |
self.vae = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae") | |
self.unet = UNetMV2DConditionModel.from_pretrained_2d(pretrained_model_path, subfolder="unet", local_crossattn=local_crossattn, **unet_from_pretrained_kwargs) | |
self.ref_unet = UNetMV2DRefModel.from_pretrained_2d(pretrained_model_path, subfolder="ref_unet", local_crossattn=local_crossattn, **unet_from_pretrained_kwargs) | |
self.text_encoder.to(device, dtype=weight_dtype) | |
self.image_encoder.to(device, dtype=weight_dtype) | |
self.vae.to(device, dtype=weight_dtype) | |
self.ref_unet.to(device, dtype=weight_dtype) | |
self.unet.to(device, dtype=weight_dtype) | |
self.vae.requires_grad_(False) | |
self.ref_unet.requires_grad_(False) | |
self.unet.requires_grad_(False) | |
self.noise_scheduler = DDIMScheduler.from_pretrained(pretrained_model_path, subfolder="scheduler-zerosnr") | |
self.validation_pipeline = CanonicalizationPipeline( | |
vae=self.vae, text_encoder=self.text_encoder, tokenizer=self.tokenizer, unet=self.unet, ref_unet=self.ref_unet,feature_extractor=self.feature_extractor,image_encoder=self.image_encoder, | |
scheduler=self.noise_scheduler | |
) | |
self.validation_pipeline.set_progress_bar_config(disable=True) | |
def canonicalize(self, image, seed): | |
return inference( | |
self.validation_pipeline, image, self.vae, self.feature_extractor, self.image_encoder, self.unet, self.ref_unet, self.tokenizer, self.text_encoder, | |
self.pretrained_model_path, self.validation, self.width_input, self.height_input, self.unet_condition_type, | |
use_noise=self.use_noise, noise_d=self.noise_d, crop=True, seed=seed, timestep=self.timestep | |
) | |
def gen(self, img_input, seed=0): | |
if np.array(img_input).shape[-1] == 4 and np.array(img_input)[..., 3].min() == 255: | |
# convert to RGB | |
img_input = img_input.convert("RGB") | |
img_output = self.canonicalize(img_input, seed) | |
max_dim = max(img_output.width, img_output.height) | |
new_image = Image.new("RGBA", (max_dim, max_dim)) | |
left = (max_dim - img_output.width) // 2 | |
top = (max_dim - img_output.height) // 2 | |
new_image.paste(img_output, (left, top)) | |
return new_image | |