import spaces from PIL import Image import io import argparse import os import sys import random import tempfile from typing import Dict, Optional, Tuple from omegaconf import OmegaConf import numpy as np import torch from pygltflib import GLTF2, Material, PbrMetallicRoughness 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 numpy as np 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, save_obj_with_mtl 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'] #### TEST #### import nvdiffrast.torch as dr import torch from typing import Tuple @spaces.GPU def _warmup(device=None): glctx = dr.RasterizeCudaContext(device=None) device = 'cuda' if device is None else device #windows workaround for https://github.com/NVlabs/nvdiffrast/issues/59 def tensor(*args, **kwargs): return torch.tensor(*args, device=device, **kwargs) pos = tensor([[[-0.8, -0.8, 0, 1], [0.8, -0.8, 0, 1], [-0.8, 0.8, 0, 1]]], dtype=torch.float32) tri = tensor([[0, 1, 2]], dtype=torch.int32) dr.rasterize(glctx, pos, tri, resolution=[256, 256]) _warmup(device) #### TEST END #### 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") @spaces.GPU def set_seed2(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, ) @spaces.GPU 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) @spaces.GPU @torch.no_grad() 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_seed2(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() @dataclass 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_image2(tensor): # 将tensor转换为numpy数组 ndarr = convert_to_numpy(tensor) # 调用save_image_numpy函数，保存图像 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 @spaces.GPU def run_multiview_infer(data, pipeline, cfg: TestConfig, num_levels=3, seed=None): pipeline.unet.enable_xformers_memory_efficient_attention() if seed is None: generator = None else: generator = torch.Generator(device=pipeline.unet.device).manual_seed(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_image2(normal) new_color = save_image2(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 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 self.results = {} # 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 infer_canonicalize_gen(img, seed) def genStage2(self, img, seed, num_levels): return infer_multiview_gen(img, seed, num_levels) def genStage3(self, img): return infer_slrm_gen(img) def genStage4(self, meshes, imgs): return infer_refine(meshes, imgs) def add_results(self, results): for k in results: self.results[k] = results[k] ############## 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 @spaces.GPU def get_distract_mask(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 return distract_mask, distract_bbox # infer_refine_sam = sam_model_registry["vit_h"](checkpoint="./ckpt/sam_vit_h_4b8939.pth").cuda() # infer_refine_generator = SamAutomaticMaskGenerator( # model=infer_refine_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, # ) infer_refine_outside_ratio = 0.20 @spaces.GPU(duration=70) def infer_refine(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 results = [] mesh_list = [] 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: # @spaces.GPU() def get_mask(): 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)) 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() return mask mask = get_mask() 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: distract_mask, distract_bbox = get_distract_mask(last_front_color, np.array(colors[0]).astype(np.float32) / 255.0) else: distract_mask = None distract_bbox = None if last_colors is None: from copy import deepcopy last_colors = deepcopy(colors) # my mesh flow weight by nearest vertexs if fixed_v is not None and fixed_f is not None and level == 1: 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() mesh_v = torch.tensor(mesh_v, dtype=torch.float32) mesh_f = torch.tensor(mesh_f) new_mesh, simp_v, simp_f = geo_refine(mesh_v, mesh_f, colors, normals, fixed_v=fixed_v, fixed_f=fixed_f) # 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.vertices.copy() 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.vertices = new_mesh_v except Exception as e: pass if fixed_v is None: fixed_v, fixed_f = simp_v, simp_f 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) mesh_list.append(new_mesh) if level == 2: new_mesh = trimesh.Trimesh(simp_v.cpu().numpy(), simp_f.cpu().numpy(), process=False) new_mesh.export(meshes[name_idx].replace('.obj', '_refined.glb')) results.append(meshes[name_idx].replace('.obj', '_refined.glb')) gltf = GLTF2().load(meshes[name_idx].replace('.obj', '_refined.glb')) for material in gltf.materials: if material.pbrMetallicRoughness: material.pbrMetallicRoughness.baseColorFactor = [1.0, 1.0, 1.0, 100.0] material.pbrMetallicRoughness.metallicFactor = 0.0 material.pbrMetallicRoughness.roughnessFactor = 1.0 gltf.save(meshes[name_idx].replace('.obj', '_refined.glb')) # save whole mesh scene = trimesh.Scene(mesh_list) scene.export(meshes[name_idx].replace('.obj', '_refined_whole.glb')) results.append(meshes[name_idx].replace('.obj', '_refined_whole.glb')) gltf = GLTF2().load(meshes[name_idx].replace('.obj', '_refined_whole.glb')) for material in gltf.materials: if material.pbrMetallicRoughness: material.pbrMetallicRoughness.baseColorFactor = [1.0, 1.0, 1.0, 100.0] material.pbrMetallicRoughness.metallicFactor = 0.0 material.pbrMetallicRoughness.roughnessFactor = 1.0 gltf.save(meshes[name_idx].replace('.obj', '_refined_whole.glb')) return results config_slrm = { 'config_path': './configs/mesh-slrm-infer.yaml' } infer_slrm_config_path = config_slrm['config_path'] infer_slrm_config = OmegaConf.load(infer_slrm_config_path) infer_slrm_config_name = os.path.basename(infer_slrm_config_path).replace('.yaml', '') infer_slrm_model_config = infer_slrm_config.model_config infer_slrm_infer_config = infer_slrm_config.infer_config infer_slrm_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") infer_slrm_model = instantiate_from_config(infer_slrm_model_config) state_dict = torch.load(infer_slrm_infer_config.model_path, map_location='cpu') infer_slrm_model.load_state_dict(state_dict, strict=False) infer_slrm_model = infer_slrm_model.to(infer_slrm_device) infer_slrm_model.init_flexicubes_geometry(infer_slrm_device, fovy=30.0, is_ortho=infer_slrm_model.is_ortho) infer_slrm_model = infer_slrm_model.eval() @spaces.GPU def infer_slrm_gen(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 = infer_slrm_make3d(imgs) return mesh_glb_fpaths[1:4] + mesh_glb_fpaths[0:1] @spaces.GPU def infer_slrm_make3d(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 = infer_slrm_model.forward_planes(images, input_cameras.float()) # get mesh mesh_glb_fpaths = [] for j in range(4): mesh_glb_fpath = infer_slrm_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 @spaces.GPU def infer_slrm_make_mesh(mesh_fpath, planes, level=None, use_texture_map=False): mesh_basename = os.path.basename(mesh_fpath).split('.')[0] mesh_dirname = os.path.dirname(mesh_fpath) with torch.no_grad(): # get mesh mesh_out = infer_slrm_model.extract_mesh( planes, use_texture_map=use_texture_map, levels=torch.tensor([level]).to(device), **infer_slrm_infer_config, ) if use_texture_map: vertices, faces, uvs, mesh_tex_idx, tex_map = mesh_out vertices = vertices[:, [1, 2, 0]] tex_map = tex_map.permute(1, 2, 0).data.cpu().numpy() if level == 2: # fill all vertex_colors with 127 tex_map = np.ones_like(tex_map) * 127 save_obj_with_mtl( vertices.data.cpu().numpy(), uvs.data.cpu().numpy(), faces.data.cpu().numpy(), mesh_tex_idx.data.cpu().numpy(), tex_map, mesh_fpath ) else: 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 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) infer_multiview_cfg = parser.parse_args() infer_multiview_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") infer_multiview_pipeline = StableUnCLIPImg2ImgPipeline.from_pretrained( infer_multiview_cfg.pretrained_path, torch_dtype=torch.float16,) if torch.cuda.is_available(): infer_multiview_pipeline.to(infer_multiview_device) print(f"Era3D Using device!!!!!!!!!!!!: {infer_multiview_device}", file=sys.stderr) infer_multiview_results = {} infer_multiview_image_transforms = [transforms.Resize(int(max(infer_multiview_cfg.height, infer_multiview_cfg.width))), transforms.CenterCrop((infer_multiview_cfg.height, infer_multiview_cfg.width)), transforms.ToTensor(), transforms.Lambda(lambda x: x * 2. - 1), ] infer_multiview_image_transforms = transforms.Compose(infer_multiview_image_transforms) prompt_embeds_path = './multiview/fixed_prompt_embeds_6view' infer_multiview_normal_text_embeds = torch.load(f'{prompt_embeds_path}/normal_embeds.pt') infer_multiview_color_text_embeds = torch.load(f'{prompt_embeds_path}/clr_embeds.pt') infer_multiview_total_views = infer_multiview_cfg.num_views @spaces.GPU def process_im(im): im = infer_multiview_image_transforms(im) return im @spaces.GPU(duration=150) def infer_multiview_gen(img, seed, num_levels): set_seed(seed) data = {} cond_im_rgb = process_im(img) cond_im_rgb = torch.stack([cond_im_rgb] * infer_multiview_total_views, dim=0) data["image_cond_rgb"] = cond_im_rgb[None, ...] data["normal_prompt_embeddings"] = infer_multiview_normal_text_embeds[None, ...] data["color_prompt_embeddings"] = infer_multiview_color_text_embeds[None, ...] results = run_multiview_infer(data, infer_multiview_pipeline, infer_multiview_cfg, num_levels=num_levels, seed=seed) return results infer_canonicalize_config = { 'config_path': './configs/canonicalization-infer.yaml', } infer_canonicalize_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # print device stderr import sys print(f"Using device!!!!!!!!!!!!: {infer_canonicalize_device}", file=sys.stderr) infer_canonicalize_config_path = infer_canonicalize_config['config_path'] infer_canonicalize_loaded_config = OmegaConf.load(infer_canonicalize_config_path) def infer_canonicalize_setup( 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, ): infer_canonicalize_width_input = width_input infer_canonicalize_height_input = height_input infer_canonicalize_timestep = timestep infer_canonicalize_use_noise = use_noise infer_canonicalize_noise_d = noise_d infer_canonicalize_validation = validation infer_canonicalize_unet_condition_type = unet_condition_type infer_canonicalize_pretrained_model_path = pretrained_model_path infer_canonicalize_local_crossattn = local_crossattn infer_canonicalize_unet_from_pretrained_kwargs = unet_from_pretrained_kwargs return infer_canonicalize_width_input, infer_canonicalize_height_input, infer_canonicalize_timestep, infer_canonicalize_use_noise, infer_canonicalize_noise_d, infer_canonicalize_validation, infer_canonicalize_unet_condition_type, infer_canonicalize_pretrained_model_path, infer_canonicalize_local_crossattn, infer_canonicalize_unet_from_pretrained_kwargs infer_canonicalize_width_input, infer_canonicalize_height_input, infer_canonicalize_timestep, infer_canonicalize_use_noise, infer_canonicalize_noise_d, infer_canonicalize_validation, infer_canonicalize_unet_condition_type, infer_canonicalize_pretrained_model_path, infer_canonicalize_local_crossattn, infer_canonicalize_unet_from_pretrained_kwargs = infer_canonicalize_setup(**infer_canonicalize_loaded_config) infer_canonicalize_tokenizer = CLIPTokenizer.from_pretrained(infer_canonicalize_pretrained_model_path, subfolder="tokenizer") infer_canonicalize_text_encoder = CLIPTextModel.from_pretrained(infer_canonicalize_pretrained_model_path, subfolder="text_encoder") infer_canonicalize_image_encoder = CLIPVisionModelWithProjection.from_pretrained(infer_canonicalize_pretrained_model_path, subfolder="image_encoder") infer_canonicalize_feature_extractor = CLIPImageProcessor() infer_canonicalize_vae = AutoencoderKL.from_pretrained(infer_canonicalize_pretrained_model_path, subfolder="vae") infer_canonicalize_unet = UNetMV2DConditionModel.from_pretrained_2d(infer_canonicalize_pretrained_model_path, subfolder="unet", local_crossattn=infer_canonicalize_local_crossattn, **infer_canonicalize_unet_from_pretrained_kwargs) infer_canonicalize_ref_unet = UNetMV2DRefModel.from_pretrained_2d(infer_canonicalize_pretrained_model_path, subfolder="ref_unet", local_crossattn=infer_canonicalize_local_crossattn, **infer_canonicalize_unet_from_pretrained_kwargs) infer_canonicalize_text_encoder.to(device, dtype=weight_dtype) infer_canonicalize_image_encoder.to(device, dtype=weight_dtype) infer_canonicalize_vae.to(device, dtype=weight_dtype) infer_canonicalize_ref_unet.to(device, dtype=weight_dtype) infer_canonicalize_unet.to(device, dtype=weight_dtype) infer_canonicalize_vae.requires_grad_(False) infer_canonicalize_ref_unet.requires_grad_(False) infer_canonicalize_unet.requires_grad_(False) infer_canonicalize_noise_scheduler = DDIMScheduler.from_pretrained(infer_canonicalize_pretrained_model_path, subfolder="scheduler-zerosnr") infer_canonicalize_validation_pipeline = CanonicalizationPipeline( vae=infer_canonicalize_vae, text_encoder=infer_canonicalize_text_encoder, tokenizer=infer_canonicalize_tokenizer, unet=infer_canonicalize_unet, ref_unet=infer_canonicalize_ref_unet,feature_extractor=infer_canonicalize_feature_extractor,image_encoder=infer_canonicalize_image_encoder, scheduler=infer_canonicalize_noise_scheduler ) infer_canonicalize_validation_pipeline.set_progress_bar_config(disable=True) @spaces.GPU def infer_canonicalize_gen(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 = inference( infer_canonicalize_validation_pipeline, img_input, infer_canonicalize_vae, infer_canonicalize_feature_extractor, infer_canonicalize_image_encoder, infer_canonicalize_unet, infer_canonicalize_ref_unet, infer_canonicalize_tokenizer, infer_canonicalize_text_encoder, infer_canonicalize_pretrained_model_path, infer_canonicalize_validation, infer_canonicalize_width_input, infer_canonicalize_height_input, infer_canonicalize_unet_condition_type, use_noise=infer_canonicalize_use_noise, noise_d=infer_canonicalize_noise_d, crop=True, seed=seed, timestep=infer_canonicalize_timestep ) 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