refactor the code

app.py

@@ -1,11 +1,11 @@
 import os
 import shutil
 import tempfile
-from functools import partial
 
 import gradio as gr
 import numpy as np
 import rembg
+import spaces
 import torch
 from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
 from einops import rearrange
@@ -16,8 +16,9 @@ from pytorch_lightning import seed_everything
 from torchvision.transforms import v2
 from tqdm import tqdm
 
-from src.utils.camera_util import FOV_to_intrinsics, get_circular_camera_poses, get_zero123plus_input_cameras
-from src.utils.infer_util import images_to_video, remove_background, resize_foreground
+from src.utils.camera_util import (FOV_to_intrinsics, get_circular_camera_poses,
+                                   get_zero123plus_input_cameras)
+from src.utils.infer_util import (remove_background, resize_foreground)
 from src.utils.mesh_util import save_glb, save_obj
 from src.utils.train_util import instantiate_from_config
 
@@ -42,50 +43,13 @@ def get_render_cameras(batch_size=1, M=120, radius=2.5, elevation=10.0, is_flexi
         cameras = cameras.unsqueeze(0).repeat(batch_size, 1, 1, 1)
     else:
         extrinsics = c2ws.flatten(-2)
-        intrinsics = FOV_to_intrinsics(50.0).unsqueeze(0).repeat(M, 1, 1).float().flatten(-2)
+        intrinsics = FOV_to_intrinsics(50.0).unsqueeze(
+            0).repeat(M, 1, 1).float().flatten(-2)
         cameras = torch.cat([extrinsics, intrinsics], dim=-1)
         cameras = cameras.unsqueeze(0).repeat(batch_size, 1, 1)
     return cameras
 
 
-def load_models(config_path):
-    config = OmegaConf.load(config_path)
-    config_name = os.path.basename(config_path).replace('.yaml', '')
-    model_config = config.model_config
-    infer_config = config.infer_config
-
-    is_flexicubes = config_name.startswith('instant-mesh')
-
-    device = torch.device('cuda')
-
-    pipeline = DiffusionPipeline.from_pretrained(
-        "sudo-ai/zero123plus-v1.2",
-        custom_pipeline="zero123plus",
-        torch_dtype=torch.float16,
-    )
-    pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
-        pipeline.scheduler.config, timestep_spacing='trailing'
-    )
-
-    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)
-
-    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(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)
-
-    model = model.to(device)
-
-    return pipeline, model, is_flexicubes, infer_config
-
-
 def check_input_image(input_image):
     if input_image is None:
         raise gr.Error("No image uploaded!")
@@ -101,27 +65,28 @@ def preprocess(input_image, do_remove_background):
     return input_image
 
 
-def generate_mvs(input_image, sample_steps, sample_seed, pipeline):
+@spaces.GPU
+def generate_mvs(input_image, sample_steps, sample_seed):
     seed_everything(sample_seed)
 
     z123_image = pipeline(
-        input_image,
-        num_inference_steps=sample_steps
-    ).images[0]
+        input_image, num_inference_steps=sample_steps).images[0]
 
     show_image = np.asarray(z123_image, dtype=np.uint8)
     show_image = torch.from_numpy(show_image)
-    show_image = rearrange(show_image, '(n h) (m w) c -> (n m) h w c', n=3, m=2)
-    show_image = rearrange(show_image, '(n m) h w c -> (n h) (m w) c', n=2, m=3)
+    show_image = rearrange(
+        show_image, '(n h) (m w) c -> (n m) h w c', n=3, m=2)
+    show_image = rearrange(
+        show_image, '(n m) h w c -> (n h) (m w) c', n=2, m=3)
     show_image = Image.fromarray(show_image.numpy())
 
     return z123_image, show_image
 
 
-def make3d(images, model, is_flexicubes, infer_config):
-    device = torch.device('cuda')
-
-    if is_flexicubes:
+@spaces.GPU
+def make3d(images):
+    global model
+    if IS_FLEXICUBES:
         model.init_flexicubes_geometry(device, use_renderer=False)
     model = model.eval()
 
@@ -129,20 +94,25 @@ def make3d(images, model, is_flexicubes, infer_config):
     images = torch.from_numpy(images).permute(2, 0, 1).contiguous().float()
     images = rearrange(images, 'c (n h) (m w) -> (n m) c h w', n=3, m=2)
 
-    input_cameras = get_zero123plus_input_cameras(batch_size=1, radius=4.0).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=2.5, 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)
+    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)
     mesh_glb_fpath = os.path.join(mesh_dirname, f"{mesh_basename}.glb")
 
     with torch.no_grad():
         planes = model.forward_planes(images, input_cameras)
-        mesh_out = model.extract_mesh(planes, use_texture_map=False, **infer_config)
+        mesh_out = model.extract_mesh(
+            planes, use_texture_map=False, **infer_config)
 
         vertices, faces, vertex_colors = mesh_out
         vertices = vertices[:, [1, 2, 0]]
@@ -150,115 +120,139 @@ def make3d(images, model, is_flexicubes, infer_config):
         save_glb(vertices, faces, vertex_colors, mesh_glb_fpath)
         save_obj(vertices, faces, vertex_colors, mesh_fpath)
 
+        print(f"Mesh saved to {mesh_fpath}")
+
     return mesh_fpath, mesh_glb_fpath
 
 
-def launch_demo(config_path):
-    cuda_path = find_cuda()
-    if cuda_path:
-        print(f"CUDA installation found at: {cuda_path}")
-    else:
-        print("CUDA installation not found")
-
-    pipeline, model, is_flexicubes, infer_config = load_models(config_path)
-
-    with gr.Blocks() as demo:
-        with gr.Row(variant="panel"):
-            with gr.Column():
-                with gr.Row():
-                    input_image = gr.Image(
-                        label="Input Image",
-                        image_mode="RGBA",
-                        sources="upload",
-                        type="pil",
-                        elem_id="content_image",
-                    )
-                    processed_image = gr.Image(
-                        label="Processed Image",
-                        image_mode="RGBA",
+# Configuration
+cuda_path = find_cuda()
+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
+infer_config = config.infer_config
+
+IS_FLEXICUBES = config_name.startswith('instant-mesh')
+device = torch.device('cuda')
+
+# Load diffusion model
+print('Loading diffusion model ...')
+pipeline = DiffusionPipeline.from_pretrained(
+    "sudo-ai/zero123plus-v1.2",
+    custom_pipeline="zero123plus",
+    torch_dtype=torch.float16,
+)
+pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+    pipeline.scheduler.config, timestep_spacing='trailing'
+)
+
+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(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)
+
+model = model.to(device)
+
+print('Loading Finished!')
+
+# Gradio UI
+with gr.Blocks() as demo:
+    with gr.Row(variant="panel"):
+        with gr.Column():
+            with gr.Row():
+                input_image = gr.Image(
+                    label="Input Image",
+                    image_mode="RGBA",
+                    sources="upload",
+                    type="pil",
+                    elem_id="content_image",
+                )
+                processed_image = gr.Image(
+                    label="Processed Image",
+                    image_mode="RGBA",
+                    type="pil",
+                    interactive=False
+                )
+            with gr.Row():
+                with gr.Group():
+                    do_remove_background = gr.Checkbox(
+                        label="Remove Background", value=True)
+                    sample_seed = gr.Number(
+                        value=42, label="Seed Value", precision=0)
+                    sample_steps = gr.Slider(
+                        label="Sample Steps", minimum=30, maximum=75, value=75, step=5)
+
+            with gr.Row():
+                submit = gr.Button(
+                    "Generate", elem_id="generate", variant="primary")
+
+            with gr.Row(variant="panel"):
+                gr.Examples(
+                    examples=[os.path.join("examples", img_name)
+                              for img_name in sorted(os.listdir("examples"))],
+                    inputs=[input_image],
+                    label="Examples",
+                    cache_examples=False,
+                    examples_per_page=16
+                )
+
+        with gr.Column():
+            with gr.Row():
+                with gr.Column():
+                    mv_show_images = gr.Image(
+                        label="Generated Multi-views",
                         type="pil",
+                        width=379,
                         interactive=False
                     )
-                with gr.Row():
-                    with gr.Group():
-                        do_remove_background = gr.Checkbox(
-                            label="Remove Background", value=True
-                        )
-                        sample_seed = gr.Number(
-                            value=42, label="Seed Value", precision=0)
-
-                        sample_steps = gr.Slider(
-                            label="Sample Steps",
-                            minimum=30,
-                            maximum=75,
-                            value=75,
-                            step=5
-                        )
-
-                with gr.Row():
-                    submit = gr.Button(
-                        "Generate", elem_id="generate", variant="primary")
-
-                with gr.Row(variant="panel"):
-                    gr.Examples(
-                        examples=[
-                            os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples"))
-                        ],
-                        inputs=[input_image],
-                        label="Examples",
-                        cache_examples=False,
-                        examples_per_page=16
-                    )
 
-            with gr.Column():
-                with gr.Row():
-                    with gr.Column():
-                        mv_show_images = gr.Image(
-                            label="Generated Multi-views",
-                            type="pil",
-                            width=379,
-                            interactive=False
-                        )
-
-                with gr.Row():
-                    with gr.Tab("OBJ"):
-                        output_model_obj = gr.Model3D(
-                            label="Output Model (OBJ Format)",
-                            interactive=False,
-                        )
-                        gr.Markdown(
-                            "Note: Downloaded .obj model will be flipped. Export .glb instead or manually flip it before usage.")
-                    with gr.Tab("GLB"):
-                        output_model_glb = gr.Model3D(
-                            label="Output Model (GLB Format)",
-                            interactive=False,
-                        )
-                        gr.Markdown(
-                            "Note: The model shown here has a darker appearance. Download to get correct results.")
-
-                with gr.Row():
+            with gr.Row():
+                with gr.Tab("OBJ"):
+                    output_model_obj = gr.Model3D(
+                        label="Output Model (OBJ Format)",
+                        interactive=False,
+                    )
+                    gr.Markdown(
+                        "Note: Downloaded .obj model will be flipped. Export .glb instead or manually flip it before usage.")
+                with gr.Tab("GLB"):
+                    output_model_glb = gr.Model3D(
+                        label="Output Model (GLB Format)",
+                        interactive=False,
+                    )
                     gr.Markdown(
-                        '''Try a different <b>seed value</b> if the result is unsatisfying (Default: 42).''')
-
-        mv_images = gr.State()
-
-        submit.click(fn=check_input_image, inputs=[input_image]).success(
-            fn=preprocess,
-            inputs=[input_image, do_remove_background],
-            outputs=[processed_image],
-        ).success(
-            fn=generate_mvs,
-            inputs=[processed_image, sample_steps, sample_seed, pipeline],
-            outputs=[mv_images, mv_show_images]
-        ).success(
-            fn=make3d,
-            inputs=[mv_images, model, is_flexicubes, infer_config],
-            outputs=[output_model_obj, output_model_glb]
-        )
-
-    demo.launch()
-
-
-if __name__ == "__main__":
-    config_path = 'configs/instant-mesh-large.yaml'
-    launch_demo(config_path)
+                        "Note: The model shown here has a darker appearance. Download to get correct results.")
+
+            with gr.Row():
+                gr.Markdown(
+                    '''Try a different <b>seed value</b> if the result is unsatisfying (Default: 42).''')
+
+    mv_images = gr.State()
+
+    submit.click(fn=check_input_image, inputs=[input_image]).success(
+        fn=preprocess,
+        inputs=[input_image, do_remove_background],
+        outputs=[processed_image],
+    ).success(
+        fn=generate_mvs,
+        inputs=[processed_image, sample_steps, sample_seed],
+        outputs=[mv_images, mv_show_images]
+    ).success(
+        fn=make3d,
+        inputs=[mv_images],
+        outputs=[output_model_obj, output_model_glb]
+    )
+
+demo.launch()