From d23bedb93e722190bfc7317407d6f6178a154caa Mon Sep 17 00:00:00 2001 From: xt4d Date: Fri, 12 Apr 2024 01:29:21 +0800 Subject: [PATCH] update app.py --- app.py | 62 ++++++++++++++++++++++++------------------------ requirements.txt | 2 +- 2 files changed, 32 insertions(+), 32 deletions(-) diff --git a/app.py b/app.py index db1f8df..9273d04 100644 --- a/app.py +++ b/app.py @@ -27,6 +27,7 @@ from src.utils.infer_util import remove_background, resize_foreground, images_to import tempfile from functools import partial +from huggingface_hub import hf_hub_download def get_render_cameras(batch_size=1, M=120, radius=2.5, elevation=10.0, is_flexicubes=False): @@ -65,7 +66,7 @@ def images_to_video(images, output_path, fps=30): seed_everything(0) -config_path = 'configs/instant-mesh-large-eval.yaml' +config_path = 'configs/instant-mesh-large.yaml' config = OmegaConf.load(config_path) config_name = os.path.basename(config_path).replace('.yaml', '') model_config = config.model_config @@ -87,15 +88,17 @@ pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config( ) # load custom white-background UNet -state_dict = torch.load(infer_config.unet_path, map_location='cpu') +unet_ckpt_path = hf_hub_download(repo_id="TencentARC/InstantMesh", filename="diffusion_pytorch_model.bin", repo_type="model") +state_dict = torch.load(unet_ckpt_path, map_location='cpu') pipeline.unet.load_state_dict(state_dict, strict=True) pipeline = pipeline.to(device) # load reconstruction model print('Loading reconstruction model ...') +model_ckpt_path = hf_hub_download(repo_id="TencentARC/InstantMesh", filename="instant_mesh_large.ckpt", repo_type="model") model = instantiate_from_config(model_config) -state_dict = torch.load(infer_config.model_path, map_location='cpu')['state_dict'] +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.') and 'source_camera' not in k} model.load_state_dict(state_dict, strict=True) @@ -115,10 +118,9 @@ def check_input_image(input_image): def preprocess(input_image, do_remove_background): rembg_session = rembg.new_session() if do_remove_background else None - - #input_image = Image.open(image_file) if do_remove_background: input_image = remove_background(input_image, rembg_session) + input_image = resize_foreground(input_image, 0.85) return input_image @@ -173,8 +175,8 @@ def make3d(images): images = torch.from_numpy(images).permute(2, 0, 1).contiguous().float() # (3, 960, 640) images = rearrange(images, 'c (n h) (m w) -> (n m) c h w', n=3, m=2) # (6, 3, 320, 320) - input_cameras = get_zero123plus_input_cameras(batch_size=1, radius=2.5).to(device) - render_cameras = get_render_cameras(batch_size=1, radius=2.5, is_flexicubes=IS_FLEXICUBES).to(device) + input_cameras = get_zero123plus_input_cameras(batch_size=1, radius=4.0).to(device) + render_cameras = get_render_cameras(batch_size=1, radius=4.0, is_flexicubes=IS_FLEXICUBES).to(device) images = images.unsqueeze(0).to(device) images = v2.functional.resize(images, (320, 320), interpolation=3, antialias=True).clamp(0, 1) @@ -223,31 +225,31 @@ def make3d(images): return video_fpath, mesh_fpath -def run_example(image_file): - - preprocessed = preprocess(image_file, False, 0.85) - mv_images, _ = generate_mvs(preprocessed, 20, 0) - video_name, mesh_fpath, planes = make3d(mv_images) - mesh_name = make_mesh(mesh_fpath, planes) - - return preprocessed, mesh_name, video_name - import gradio as gr -HEADER = ''' -

-Official 🤗 Gradio demo for - -InstantMesh: Efficient 3D Mesh Generation from a Single Image with Sparse-view Large Reconstruction Models -. -

-
-* If the output is unsatisfying, try to use a different seed. +_HEADER_ = ''' +

Official 🤗 Gradio Demo

InstantMesh: Efficient 3D Mesh Generation from a Single Image with Sparse-view Large Reconstruction Models

''' +_LINKS_ = ''' +

Code is available at GitHub

+

Report is available at ArXiv

+''' + +_CITE_ = r""" +```bibtex +@article{xu2024instantmesh, + title={InstantMesh: Efficient 3D Mesh Generation from a Single Image with Sparse-view Large Reconstruction Models}, + author={Xu, Jiale and Cheng, Weihao and Gao, Yiming and Wang, Xintao and Gao, Shenghua and Shan, Ying}, + journal={arXiv preprint arXiv:2404.07191}, + year={2024} +} +``` +""" + with gr.Blocks() as demo: - gr.Markdown(HEADER) + gr.Markdown(_HEADER_) with gr.Row(variant="panel"): with gr.Column(): with gr.Row(): @@ -273,7 +275,7 @@ with gr.Blocks() as demo: do_remove_background = gr.Checkbox( label="Remove Background", value=True ) - sample_seed = gr.Number(value=42, label="Seed", precision=0) + sample_seed = gr.Number(value=42, label="Seed (Try a different value if the result is unsatisfying)", precision=0) sample_steps = gr.Slider( label="Sample Steps", @@ -292,9 +294,6 @@ with gr.Blocks() as demo: os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples")) ], inputs=[input_image], - # outputs=[processed_image, output_model_obj, output_video], - # fn=partial(run_example), - # cache_examples=True, label="Examples", examples_per_page=20 ) @@ -325,7 +324,8 @@ with gr.Blocks() as demo: width=768, interactive=False, ) - + gr.Markdown(_LINKS_) + gr.Markdown(_CITE_) mv_images = gr.State() submit.click(fn=check_input_image, inputs=[input_image]).success( diff --git a/requirements.txt b/requirements.txt index ab2130e..aa68bbd 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ pytorch-lightning==2.1.2 -gradio +gradio==3.41.2 huggingface-hub einops omegaconf