import os from einops import rearrange from omegaconf import OmegaConf import torch import numpy as np import trimesh import torchvision import torch.nn.functional as F from PIL import Image from torchvision import transforms from torchvision.transforms import v2 from transformers import AutoProcessor, AutoModelForCausalLM import rembg from diffusers import FluxPipeline, FluxControlNetImg2ImgPipeline from diffusers.models.controlnet_flux import FluxControlNetModel, FluxMultiControlNetModel from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler, HeunDiscreteScheduler from pytorch_lightning import seed_everything import os from models.ISOMER.reconstruction_func import reconstruction from models.ISOMER.projection_func import projection from models.lrm.utils.infer_util import remove_background, resize_foreground, save_video from models.lrm.utils.mesh_util import save_obj, save_obj_with_mtl from models.lrm.utils.render_utils import rotate_x, rotate_y from models.lrm.utils.train_util import instantiate_from_config from models.lrm.utils.camera_util import get_zero123plus_input_cameras, get_custom_zero123plus_input_cameras, get_flux_input_cameras from utils.tool import NormalTransfer, get_render_cameras_frames, load_mipmap from utils.tool import get_background, get_render_cameras_video, render_frames, mask_fix device = "cuda" resolution = 512 save_dir = "./outputs" zero123plus_diffusion_steps = 75 normal_transfer = NormalTransfer() rembg_session = rembg.new_session() isomer_azimuths = torch.from_numpy(np.array([270, 0, 90, 180])).to(device) isomer_elevations = torch.from_numpy(np.array([5, 5, 5, 5])).to(device) isomer_radius = 4.1 isomer_geo_weights = torch.from_numpy(np.array([1, 0.9, 1, 0.9])).float().to(device) isomer_color_weights = torch.from_numpy(np.array([1, 0.5, 1, 0.5])).float().to(device) # seed_everything(42) # model initialization and loading # flux print('==> Loading Flux model ...') flux_base_model_pth = "/hpc2hdd/JH_DATA/share/yingcongchen/PrivateShareGroup/yingcongchen_datasets/model_checkpoint/models--black-forest-labs--FLUX.1-dev" flux_controlnet = FluxControlNetModel.from_pretrained("/hpc2hdd/JH_DATA/share/yingcongchen/PrivateShareGroup/yingcongchen_datasets/model_checkpoint/flux_controlnets/FLUX.1-dev-ControlNet-Union-Pro") flux_pipe = FluxControlNetImg2ImgPipeline.from_pretrained(flux_base_model_pth, controlnet=[flux_controlnet], torch_dtype=torch.bfloat16).to(device=device, dtype=torch.bfloat16) flux_pipe.load_lora_weights('./checkpoint/flux_lora/rgb_normal_large.safetensors') flux_pipe.to(device=device, dtype=torch.bfloat16) generator = torch.Generator(device=device).manual_seed(0) # lrm print('==> Loading LRM model ...') config = OmegaConf.load("./models/lrm/config/PRM_inference.yaml") model_config = config.model_config infer_config = config.infer_config model = instantiate_from_config(model_config) model_ckpt_path = "./checkpoint/lrm/final_ckpt.ckpt" state_dict = torch.load(model_ckpt_path, map_location='cpu')['state_dict'] state_dict = {k[14:]: v for k, v in state_dict.items() if k.startswith('lrm_generator.')} model.load_state_dict(state_dict, strict=True) model = model.to(device) model.init_flexicubes_geometry(device, fovy=50.0) model = model.eval() # zero123++ print('==> Loading diffusion model ...') zero123plus_pipeline = DiffusionPipeline.from_pretrained( "sudo-ai/zero123plus-v1.2", custom_pipeline="./models/zero123plus", torch_dtype=torch.float16, ) zero123plus_pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config( zero123plus_pipeline.scheduler.config, timestep_spacing='trailing' ) unet_ckpt_path = "./checkpoint/zero123++/flexgen_19w.ckpt" state_dict = torch.load(unet_ckpt_path, map_location='cpu')['state_dict'] state_dict = {k[10:]: v for k, v in state_dict.items() if k.startswith('unet.unet.')} zero123plus_pipeline.unet.load_state_dict(state_dict, strict=True) zero123plus_pipeline = zero123plus_pipeline.to(device) # unet_ckpt_path = "checkpoint/zero123++/diffusion_pytorch_model.bin" # state_dict = torch.load(unet_ckpt_path, map_location='cpu') # zero123plus_pipeline.unet.load_state_dict(state_dict, strict=True) # zero123plus_pipeline = zero123plus_pipeline.to(device) # florence caption_model = AutoModelForCausalLM.from_pretrained( "/hpc2hdd/home/jlin695/.cache/huggingface/hub/models--multimodalart--Florence-2-large-no-flash-attn/snapshots/8db3793cf5b453b2ccfb3a4f613b403b2e6b7ca2", torch_dtype=torch.bfloat16, trust_remote_code=True, ).to(device) caption_processor = AutoProcessor.from_pretrained("/hpc2hdd/home/jlin695/.cache/huggingface/hub/models--multimodalart--Florence-2-large-no-flash-attn/snapshots/8db3793cf5b453b2ccfb3a4f613b403b2e6b7ca2", trust_remote_code=True) # Flux multi-view generation def multi_view_rgb_normal_generation_with_controlnet(prompt, image, strength=1.0, control_image=[], control_mode=[], control_guidance_start=None, control_guidance_end=None, controlnet_conditioning_scale=None, lora_scale=1.0 ): control_mode_dict = { 'canny': 0, 'tile': 1, 'depth': 2, 'blur': 3, 'pose': 4, 'gray': 5, 'lq': 6, } # for https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union only hparam_dict = { 'prompt': prompt, 'image': image, 'strength': strength, 'num_inference_steps': 30, 'guidance_scale': 3.5, 'num_images_per_prompt': 1, 'width': resolution*4, 'height': resolution*2, 'output_type': 'np', 'generator': generator, 'joint_attention_kwargs': {"scale": lora_scale} } # append controlnet hparams if len(control_image) > 0: assert len(control_mode) == len(control_image) # the count of image should be the same as control mode ctrl_hparams = { 'control_mode': [control_mode_dict[mode_] for mode_ in control_mode], 'control_image': control_image, 'control_guidance_start': control_guidance_start or [0.0 for i in range(len(control_image))], 'control_guidance_end': control_guidance_end or [1.0 for i in range(len(control_image))], 'controlnet_conditioning_scale': controlnet_conditioning_scale or [1.0 for i in range(len(control_image))], } hparam_dict.update(ctrl_hparams) # generate multi-view images with torch.no_grad(): image = flux_pipe( **hparam_dict ).images return image # captioning def run_captioning(image): device = "cuda" if torch.cuda.is_available() else "cpu" torch_dtype = torch.bfloat16 if isinstance(image, str): # If image is a file path image = Image.open(image).convert("RGB") prompt = "" inputs = caption_processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype) # print(f"inputs {inputs}") generated_ids = caption_model.generate( input_ids=inputs["input_ids"], pixel_values=inputs["pixel_values"], max_new_tokens=1024, num_beams=3 ) generated_text = caption_processor.batch_decode(generated_ids, skip_special_tokens=False)[0] parsed_answer = caption_processor.post_process_generation( generated_text, task=prompt, image_size=(image.width, image.height) ) # print(f"parsed_answer = {parsed_answer}") caption_text = parsed_answer[""].replace("The image is ", "") return caption_text # zero123++ multi-view generation def multi_view_rgb_generation(cond_img): # generate multi-view images with torch.no_grad(): output_image = zero123plus_pipeline( cond_img, num_inference_steps=zero123plus_diffusion_steps, width=resolution*2, height=resolution*2, ).images[0] return output_image # lrm reconstructions def lrm_reconstructions(image, input_cameras, save_path=None, name="temp", export_texmap=False, if_save_video=False, render_azimuths=None, render_elevations=None, render_radius=None, render_fov=30): images = image.unsqueeze(0).to(device) images = v2.functional.resize(images, 512, interpolation=3, antialias=True).clamp(0, 1) # breakpoint() with torch.no_grad(): # get triplane planes = model.forward_planes(images, input_cameras) mesh_path_idx = os.path.join(save_path, f'{name}.obj') mesh_out = model.extract_mesh( planes, use_texture_map=export_texmap, **infer_config, ) if export_texmap: vertices, faces, uvs, mesh_tex_idx, tex_map = mesh_out save_obj_with_mtl( vertices.data.cpu().numpy(), uvs.data.cpu().numpy(), faces.data.cpu().numpy(), mesh_tex_idx.data.cpu().numpy(), tex_map.permute(1, 2, 0).data.cpu().numpy(), mesh_path_idx, ) else: vertices, faces, vertex_colors = mesh_out save_obj(vertices, faces, vertex_colors, mesh_path_idx) print(f"Mesh saved to {mesh_path_idx}") render_size = 512 if if_save_video: video_path_idx = os.path.join(save_path, f'{name}.mp4') render_size = infer_config.render_resolution ENV = load_mipmap("models/lrm/env_mipmap/6") materials = (0.0,0.9) all_mv, all_mvp, all_campos = get_render_cameras_video( batch_size=1, M=240, radius=4.5, elevation=(90, 60.0), is_flexicubes=True, fov=30 ) frames, albedos, pbr_spec_lights, pbr_diffuse_lights, normals, alphas = render_frames( model, planes, render_cameras=all_mvp, camera_pos=all_campos, env=ENV, materials=materials, render_size=render_size, chunk_size=20, is_flexicubes=True, ) normals = (torch.nn.functional.normalize(normals) + 1) / 2 normals = normals * alphas + (1-alphas) all_frames = torch.cat([frames, albedos, pbr_spec_lights, pbr_diffuse_lights, normals], dim=3) # breakpoint() save_video( all_frames, video_path_idx, fps=30, ) print(f"Video saved to {video_path_idx}") if render_azimuths is not None and render_elevations is not None and render_radius is not None: render_size = infer_config.render_resolution ENV = load_mipmap("models/lrm/env_mipmap/6") materials = (0.0,0.9) all_mv, all_mvp, all_campos, identity_mv = get_render_cameras_frames( batch_size=1, radius=render_radius, azimuths=render_azimuths, elevations=render_elevations, fov=30 ) frames, albedos, pbr_spec_lights, pbr_diffuse_lights, normals, alphas = render_frames( model, planes, render_cameras=all_mvp, camera_pos=all_campos, env=ENV, materials=materials, render_size=render_size, render_mv = all_mv, local_normal=True, identity_mv=identity_mv, ) else: normals = None frames = None albedos = None return vertices, faces, normals, frames, albedos def transform_normal(input_normal, azimuths_deg, elevations_deg, radius=4.5, is_global_to_local=False): """ input_normal: in range [-1, 1], shape (b c h w) """ input_normal = input_normal.permute(0, 2, 3, 1).cpu() azimuths_deg = np.array(azimuths_deg) elevations_deg = np.array(elevations_deg) if is_global_to_local: local_normal = normal_transfer.trans_global_2_local(input_normal, azimuths_deg, elevations_deg) return local_normal.permute(0, 3, 1, 2) else: global_normal = normal_transfer.trans_local_2_global(input_normal, azimuths_deg, elevations_deg, radius=radius, for_lotus=False) global_normal[..., 0] *= -1 return global_normal.permute(0, 3, 1, 2) def local_normal_global_transform(local_normal_images,azimuths_deg,elevations_deg): if local_normal_images.min() >= 0: local_normal = local_normal_images.float() * 2 - 1 else: local_normal = local_normal_images.float() global_normal = normal_transfer.trans_local_2_global(local_normal, azimuths_deg, elevations_deg, radius=4.5, for_lotus=False) global_normal[...,0] *= -1 global_normal = (global_normal + 1) / 2 global_normal = global_normal.permute(0, 3, 1, 2) return global_normal def main(): image_pth = "examples/蓝色小怪物.webp" save_dir_path = os.path.join(save_dir, image_pth.split("/")[-1].split(".")[0]) os.makedirs(save_dir_path, exist_ok=True) input_image = Image.open(image_pth) # if not args.no_rembg: input_image = remove_background(input_image, rembg_session) input_image = resize_foreground(input_image, 0.85) # generate caption image_caption = run_captioning(image_pth) # generate multi-view images output_image = multi_view_rgb_generation(input_image) # lrm reconstructions rgb_multi_view = np.asarray(output_image, dtype=np.float32) / 255.0 rgb_multi_view = torch.from_numpy(rgb_multi_view).squeeze(0).permute(2, 0, 1).contiguous().float() # (3, 1024, 2048) rgb_multi_view = rearrange(rgb_multi_view, 'c (n h) (m w) -> (n m) c h w', n=2, m=2) # (8, 3, 512, 512) input_cameras = get_custom_zero123plus_input_cameras(batch_size=1, radius=3.5, fov=30).to(device) vertices, faces, lrm_multi_view_normals, lrm_multi_view_rgb, lrm_multi_view_albedo = \ lrm_reconstructions(rgb_multi_view, input_cameras, save_path=save_dir_path, name='lrm', export_texmap=False, if_save_video=False, render_azimuths=isomer_azimuths, render_elevations=isomer_elevations, render_radius=isomer_radius, render_fov=30) vertices = torch.from_numpy(vertices).to(device) faces = torch.from_numpy(faces).to(device) vertices = vertices @ rotate_x(np.pi / 2, device=vertices.device)[:3, :3] vertices = vertices @ rotate_y(np.pi / 2, device=vertices.device)[:3, :3] # lrm_3D_bundle_image = torchvision.utils.make_grid(torch.cat([lrm_multi_view_rgb.cpu(), (lrm_multi_view_normals.cpu() + 1) / 2], dim=0), nrow=4, padding=0).unsqueeze(0) # range [0, 1] lrm_3D_bundle_image = torchvision.utils.make_grid(torch.cat([rgb_multi_view[[3,0,1,2]].cpu(), (lrm_multi_view_normals.cpu() + 1) / 2], dim=0), nrow=4, padding=0).unsqueeze(0) # range [0, 1] # rgb_multi_view[[3,0,1,2]] : (B,3,H,W) # lrm_multi_view_normals : (B,3,H,W) # combined_images = 0.5 * rgb_multi_view[[3,0,1,2]].cpu() + 0.5 * (lrm_multi_view_normals.cpu() + 1) / 2 # torchvision.utils.save_image(combined_images, os.path.join("debug_output", 'combined.png')) # breakpoint() # Use the low-quality controlnet by default, feel free to try the others control_image = [lrm_3D_bundle_image * 2 - 1] control_mode = ['tile'] control_guidance_start = [0.0] control_guidance_end = [0.3] controlnet_conditioning_scale = [0.8] flux_pipe.controlnet = FluxMultiControlNetModel([flux_controlnet for _ in control_mode]) # breakpoint() rgb_normal_grid = multi_view_rgb_normal_generation_with_controlnet( prompt= ' '.join(['A grid of 2x4 multi-view image, elevation 5. White background.', image_caption]), image=lrm_3D_bundle_image, strength=0.6, control_image=control_image, control_mode=control_mode, control_guidance_start=control_guidance_start, control_guidance_end=control_guidance_end, controlnet_conditioning_scale=controlnet_conditioning_scale, lora_scale=1.0 ) # noted that rgb_normal_grid is a (b, h, w, c) numpy array rgb_normal_grid = torch.from_numpy(rgb_normal_grid).contiguous().float() rgb_normal_grid = rearrange(rgb_normal_grid.squeeze(0), '(n h) (m w) c-> (n m) c h w', n=2, m=4) # (8, 3, 512, 512) rgb_multi_view = rgb_normal_grid[:4, :3, :, :].cuda() normal_multi_view = rgb_normal_grid[4:, :3, :, :].cuda() multi_view_mask = get_background(normal_multi_view).cuda() rgb_multi_view = rgb_multi_view * multi_view_mask + (1-multi_view_mask) # local normal to global normal global_normal = local_normal_global_transform(normal_multi_view.permute(0, 2, 3, 1).cpu(), isomer_azimuths, isomer_elevations).cuda() global_normal = global_normal * multi_view_mask + (1-multi_view_mask) global_normal = global_normal.permute(0,2,3,1) multi_view_mask = multi_view_mask.squeeze(1) rgb_multi_view = rgb_multi_view.permute(0,2,3,1) # global_normal: B,H,W,3 # multi_view_mask: B,H,W # rgb_multi_view: B,H,W,3 meshes = reconstruction( normal_pils=global_normal, masks=multi_view_mask, weights=isomer_geo_weights, fov=30, radius=isomer_radius, camera_angles_azi=isomer_azimuths, camera_angles_ele=isomer_elevations, expansion_weight_stage1=0.1, init_type="file", init_verts=vertices, init_faces=faces, stage1_steps=0, stage2_steps=50, start_edge_len_stage1=0.1, end_edge_len_stage1=0.02, start_edge_len_stage2=0.02, end_edge_len_stage2=0.005, ) save_glb_addr = projection( meshes=meshes, masks=multi_view_mask, images=rgb_multi_view, azimuths=isomer_azimuths, elevations=isomer_elevations, weights=isomer_color_weights, fov=30, radius=isomer_radius, save_dir=f"{save_dir_path}/ISOMER/", ) print(f'saved to {save_glb_addr}') if __name__ == '__main__': main()