Spaces:
Sleeping
Sleeping
| import os | |
| import imageio | |
| import numpy as np | |
| os.system("bash install.sh") | |
| from omegaconf import OmegaConf | |
| import tqdm | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import torchvision.transforms.functional as TF | |
| import rembg | |
| import gradio as gr | |
| from gradio_litmodel3d import LitModel3D | |
| from dva.io import load_from_config | |
| from dva.ray_marcher import RayMarcher | |
| from dva.visualize import visualize_primvolume, visualize_video_primvolume | |
| from inference import remove_background, resize_foreground, extract_texmesh | |
| from models.diffusion import create_diffusion | |
| from huggingface_hub import hf_hub_download | |
| ckpt_path = hf_hub_download(repo_id="frozenburning/3DTopia-XL", filename="model_sview_dit_fp16.pt") | |
| vae_ckpt_path = hf_hub_download(repo_id="frozenburning/3DTopia-XL", filename="model_vae_fp16.pt") | |
| GRADIO_PRIM_VIDEO_PATH = 'prim.mp4' | |
| GRADIO_RGB_VIDEO_PATH = 'rgb.mp4' | |
| GRADIO_MAT_VIDEO_PATH = 'mat.mp4' | |
| GRADIO_GLB_PATH = 'pbr_mesh.glb' | |
| CONFIG_PATH = "./configs/inference_dit.yml" | |
| config = OmegaConf.load(CONFIG_PATH) | |
| config.checkpoint_path = ckpt_path | |
| config.model.vae_checkpoint_path = vae_ckpt_path | |
| # model | |
| model = load_from_config(config.model.generator) | |
| state_dict = torch.load(config.checkpoint_path, map_location='cpu') | |
| model.load_state_dict(state_dict['ema']) | |
| vae = load_from_config(config.model.vae) | |
| vae_state_dict = torch.load(config.model.vae_checkpoint_path, map_location='cpu') | |
| vae.load_state_dict(vae_state_dict['model_state_dict']) | |
| conditioner = load_from_config(config.model.conditioner) | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| vae = vae.to(device) | |
| conditioner = conditioner.to(device) | |
| model = model.to(device) | |
| model.eval() | |
| amp = True | |
| precision_dtype = torch.float16 | |
| rm = RayMarcher( | |
| 256, | |
| 256, | |
| **config.rm, | |
| ).to(device) | |
| perchannel_norm = False | |
| if "latent_mean" in config.model: | |
| latent_mean = torch.Tensor(config.model.latent_mean)[None, None, :].to(device) | |
| latent_std = torch.Tensor(config.model.latent_std)[None, None, :].to(device) | |
| assert latent_mean.shape[-1] == config.model.generator.in_channels | |
| perchannel_norm = True | |
| latent_nf = config.model.latent_nf | |
| config.diffusion.pop("timestep_respacing") | |
| config.model.pop("vae") | |
| config.model.pop("vae_checkpoint_path") | |
| config.model.pop("conditioner") | |
| config.model.pop("generator") | |
| config.model.pop("latent_nf") | |
| config.model.pop("latent_mean") | |
| config.model.pop("latent_std") | |
| model_primx = load_from_config(config.model) | |
| # load rembg | |
| rembg_session = rembg.new_session() | |
| # background removal function | |
| def background_remove_process(input_image): | |
| input_image = remove_background(input_image, rembg_session) | |
| input_image = resize_foreground(input_image, 0.85) | |
| input_cond_preview_pil = input_image | |
| raw_image = np.array(input_image) | |
| mask = (raw_image[..., -1][..., None] > 0) * 1 | |
| raw_image = raw_image[..., :3] * mask | |
| input_cond = torch.from_numpy(np.array(raw_image)[None, ...]).to(device) | |
| return gr.update(interactive=True), input_cond, input_cond_preview_pil | |
| # process function | |
| def process(input_cond, input_num_steps, input_seed=42, input_cfg=6.0): | |
| # seed | |
| torch.manual_seed(input_seed) | |
| os.makedirs(config.output_dir, exist_ok=True) | |
| output_rgb_video_path = os.path.join(config.output_dir, GRADIO_RGB_VIDEO_PATH) | |
| output_prim_video_path = os.path.join(config.output_dir, GRADIO_PRIM_VIDEO_PATH) | |
| output_mat_video_path = os.path.join(config.output_dir, GRADIO_MAT_VIDEO_PATH) | |
| respacing = "ddim{}".format(input_num_steps) | |
| diffusion = create_diffusion(timestep_respacing=respacing, **config.diffusion) | |
| sample_fn = diffusion.ddim_sample_loop_progressive | |
| fwd_fn = model.forward_with_cfg | |
| # text-conditioned | |
| if input_cond is None: | |
| raise NotImplementedError | |
| with torch.no_grad(): | |
| latent = torch.randn(1, config.model.num_prims, 1, 4, 4, 4) | |
| batch = {} | |
| inf_bs = 1 | |
| inf_x = torch.randn(inf_bs, config.model.num_prims, 68).to(device) | |
| y = conditioner.encoder(input_cond) | |
| model_kwargs = dict(y=y[:inf_bs, ...], precision_dtype=precision_dtype, enable_amp=amp) | |
| if input_cfg >= 0: | |
| model_kwargs['cfg_scale'] = input_cfg | |
| for samples in sample_fn(fwd_fn, inf_x.shape, inf_x, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device): | |
| final_samples = samples | |
| recon_param = final_samples["sample"].reshape(inf_bs, config.model.num_prims, -1) | |
| if perchannel_norm: | |
| recon_param = recon_param / latent_nf * latent_std + latent_mean | |
| recon_srt_param = recon_param[:, :, 0:4] | |
| recon_feat_param = recon_param[:, :, 4:] # [8, 2048, 64] | |
| recon_feat_param_list = [] | |
| # one-by-one to avoid oom | |
| for inf_bidx in range(inf_bs): | |
| if not perchannel_norm: | |
| decoded = vae.decode(recon_feat_param[inf_bidx, ...].reshape(1*config.model.num_prims, *latent.shape[-4:]) / latent_nf) | |
| else: | |
| decoded = vae.decode(recon_feat_param[inf_bidx, ...].reshape(1*config.model.num_prims, *latent.shape[-4:])) | |
| recon_feat_param_list.append(decoded.detach()) | |
| recon_feat_param = torch.concat(recon_feat_param_list, dim=0) | |
| # invert normalization | |
| if not perchannel_norm: | |
| recon_srt_param[:, :, 0:1] = (recon_srt_param[:, :, 0:1] / 10) + 0.05 | |
| recon_feat_param[:, 0:1, ...] /= 5. | |
| recon_feat_param[:, 1:, ...] = (recon_feat_param[:, 1:, ...] + 1) / 2. | |
| recon_feat_param = recon_feat_param.reshape(inf_bs, config.model.num_prims, -1) | |
| recon_param = torch.concat([recon_srt_param, recon_feat_param], dim=-1) | |
| visualize_video_primvolume(config.output_dir, batch, recon_param, 15, rm, device) | |
| prim_params = {'srt_param': recon_srt_param[0].detach().cpu(), 'feat_param': recon_feat_param[0].detach().cpu()} | |
| return output_rgb_video_path, output_prim_video_path, output_mat_video_path, prim_params | |
| def export_mesh(prim_params, uv_unwrap="Faster", remesh="No", mc_resolution=256): | |
| # exporting GLB mesh | |
| output_glb_path = os.path.join(config.output_dir, GRADIO_GLB_PATH) | |
| if remesh == "No": | |
| config.inference.remesh = False | |
| elif remesh == "Yes": | |
| config.inference.remesh = True | |
| if uv_unwrap == "Faster": | |
| config.inference.fast_unwrap = True | |
| elif uv_unwrap == "Better": | |
| config.inference.fast_unwrap = False | |
| config.inference.mc_resolution = mc_resolution | |
| config.inference.batch_size = 8192 | |
| model_primx.load_state_dict(prim_params) | |
| model_primx.to(device) | |
| model_primx.eval() | |
| with torch.no_grad(): | |
| model_primx.srt_param[:, 1:4] *= 0.85 | |
| extract_texmesh(config.inference, model_primx, config.output_dir, device) | |
| return output_glb_path, gr.update(visible=True), gr.update(interactive=True), gr.update(value="assets/hdri/metro_noord_1k.hdr") | |
| # gradio UI | |
| _TITLE = '''3DTopia-XL: Scaling High-quality 3D Asset Generation via Primitive Diffusion''' | |
| _DESCRIPTION = ''' | |
| <div> | |
| <a style="display:inline-block" href="https://frozenburning.github.io/projects/3DTopia-XL/"><img src='https://img.shields.io/badge/public_website-8A2BE2'></a> | |
| <a style="display:inline-block; margin-left: .5em" href="https://github.com/3DTopia/3DTopia-XL"><img src='https://img.shields.io/github/stars/3DTopia/3DTopia-XL?style=social'/></a> | |
| </div> | |
| * Now we offer 1) single image conditioned model, we will release 2) multiview images conditioned model and 3) pure text conditioned model in the future! | |
| * If you find the output unsatisfying, try using different seeds! | |
| ''' | |
| _DEV_DES = ''' | |
| * Please refer to our repo for instructions on running gradio demo [locally](https://github.com/3DTopia/3DTopia-XL?tab=readme-ov-file#gradio-demo) or [CLI test](https://github.com/3DTopia/3DTopia-XL?tab=readme-ov-file#cli-test) | |
| ''' | |
| block = gr.Blocks(title=_TITLE).queue() | |
| with block: | |
| current_fg_state = gr.State() | |
| prim_param_state = gr.State() | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown('# ' + _TITLE) | |
| gr.Markdown(_DESCRIPTION) | |
| with gr.Accordion("For Developers", open=False): | |
| gr.Markdown(_DEV_DES) | |
| with gr.Row(variant='panel'): | |
| with gr.Column(scale=1): | |
| with gr.Row(): | |
| # input image | |
| input_image = gr.Image(label="image", type='pil') | |
| # background removal | |
| removal_previewer = gr.Image(label="Background Removal Preview", type='pil', interactive=False) | |
| with gr.Row(): | |
| # inference steps | |
| input_num_steps = gr.Radio(choices=[25, 50, 100, 200], label="DDIM steps", value=25, info="Larger for robustness but slower.") | |
| # random seed | |
| input_cfg = gr.Slider(label="CFG scale", minimum=0, maximum=15, step=0.5, value=6, info="Typically CFG in a range of 4-7") | |
| # random seed | |
| input_seed = gr.Slider(label="random seed", minimum=0, maximum=10000, step=1, value=42, info="Try different seed if the result is not satisfying as this is a generative model!") | |
| with gr.Row(): | |
| input_mc_resolution = gr.Radio(choices=[128, 256], label="MC Resolution", value=128, info="Cube resolution for mesh extraction. Larger for better quality but slower.") | |
| input_remesh = gr.Radio(choices=["No", "Yes"], label="Remesh", value="No", info="Remesh or not?") | |
| input_unwrap = gr.Radio(choices=["Faster", "Better"], label="UV Unwrap", value="Better", info="UV unwrapping algorithm. Trade-off between quality and speed.") | |
| # gen button | |
| with gr.Row(): | |
| button_gen = gr.Button(value="Generate", interactive=False) | |
| export_glb_btn = gr.Button(value="Export Current GLB", interactive=False) | |
| with gr.Column(scale=1): | |
| with gr.Row(): | |
| # final video results | |
| output_rgb_video = gr.Video(label="RGB") | |
| output_prim_video = gr.Video(label="Primitives") | |
| output_mat_video = gr.Video(label="Material") | |
| with gr.Row(): | |
| # glb file | |
| output_glb = LitModel3D( | |
| label="3D GLB Model", | |
| visible=True, | |
| clear_color=[0.0, 0.0, 0.0, 0.0], | |
| camera_position=(90, None, None), | |
| tonemapping="aces", | |
| contrast=1.0, | |
| scale=1.0, | |
| ) | |
| with gr.Column(visible=False, scale=1.0) as hdr_row: | |
| gr.Markdown("""## HDR Environment Map | |
| Select / Upload an HDR environment map to relight the 3D model. | |
| """) | |
| with gr.Row(): | |
| example_hdris = [ | |
| os.path.join("assets/hdri", f) | |
| for f in os.listdir("assets/hdri") | |
| ] | |
| hdr_illumination_file = gr.File( | |
| label="HDR Envmap", file_types=[".hdr"], file_count="single" | |
| ) | |
| hdr_illumination_example = gr.Examples( | |
| examples=example_hdris, | |
| inputs=hdr_illumination_file, | |
| ) | |
| hdr_illumination_file.change( | |
| lambda x: gr.update(env_map=x.name if x is not None else None), | |
| inputs=hdr_illumination_file, | |
| outputs=[output_glb], | |
| ) | |
| input_image.change(background_remove_process, inputs=[input_image], outputs=[button_gen, current_fg_state, removal_previewer]) | |
| button_gen.click(process, inputs=[current_fg_state, input_num_steps, input_seed, input_cfg], outputs=[output_rgb_video, output_prim_video, output_mat_video, prim_param_state]) | |
| prim_param_state.change(export_mesh, inputs=[prim_param_state, input_unwrap, input_remesh, input_mc_resolution], outputs=[output_glb, hdr_row, export_glb_btn, hdr_illumination_file]) | |
| export_glb_btn.click(export_mesh, inputs=[prim_param_state, input_unwrap, input_remesh, input_mc_resolution], outputs=[output_glb, hdr_row, export_glb_btn, hdr_illumination_file]) | |
| gr.Examples( | |
| examples=[ | |
| os.path.join("assets/examples", f) | |
| for f in os.listdir("assets/examples") | |
| ], | |
| inputs=[input_image], | |
| outputs=[output_rgb_video, output_prim_video, output_mat_video, prim_param_state], | |
| fn=lambda x: process(input_image=x), | |
| cache_examples=False, | |
| label='Single Image to 3D PBR Asset' | |
| ) | |
| block.launch(server_name="0.0.0.0", share=True) |