Refactor app.py by moving model initialization and argument parsing to the main execution block, enhancing clarity and organization. Correct typos in comments and variable names for improved readability. Update Gradio component configurations for better user experience.

app.py

@@ -1,4 +1,3 @@
-# Fixed version with proper error handling and compatibility
 import spaces
 import argparse
 import numpy as np
@@ -20,71 +19,10 @@ import argparse
 from model import CRM
 from inference import generate3d
 
-# Move model initialization into a function that will be called by workers
-def init_model():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--stage1_config",
-        type=str,
-        default="configs/nf7_v3_SNR_rd_size_stroke.yaml",
-        help="config for stage1",
-    )
-    parser.add_argument(
-        "--stage2_config",
-        type=str,
-        default="configs/stage2-v2-snr.yaml",
-        help="config for stage2",
-    )
-    parser.add_argument("--device", type=str, default="cuda")
-    args = parser.parse_args(args=[])  # Fix: provide empty args list
-
-    # Download model files
-    crm_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="CRM.pth")
-    specs = json.load(open("configs/specs_objaverse_total.json"))
-    model = CRM(specs)
-    model.load_state_dict(torch.load(crm_path, map_location="cpu"), strict=False)
-    model = model.to(args.device)
-
-    # Load configs
-    stage1_config = OmegaConf.load(args.stage1_config).config
-    stage2_config = OmegaConf.load(args.stage2_config).config
-    stage2_sampler_config = stage2_config.sampler
-    stage1_sampler_config = stage1_config.sampler
-
-    stage1_model_config = stage1_config.models
-    stage2_model_config = stage2_config.models
-
-    xyz_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="ccm-diffusion.pth")
-    pixel_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="pixel-diffusion.pth")
-    stage1_model_config.resume = pixel_path
-    stage2_model_config.resume = xyz_path
-
-    pipeline = TwoStagePipeline(
-        stage1_model_config,
-        stage2_model_config,
-        stage1_sampler_config,
-        stage2_sampler_config,
-        device=args.device,
-        dtype=torch.float32
-    )
-    
-    return model, pipeline, args
-
-# Global variables to store model and pipeline
-model = None
 pipeline = None
-args = None
-
-@spaces.GPU
-def get_model():
-    """Lazy initialization of model and pipeline"""
-    global model, pipeline, args
-    if model is None or pipeline is None:
-        model, pipeline, args = init_model()
-    return model, pipeline
-
 rembg_session = rembg.new_session()
 
+
 def expand_to_square(image, bg_color=(0, 0, 0, 0)):
     # expand image to 1:1
     width, height = image.size
@@ -97,10 +35,9 @@ def expand_to_square(image, bg_color=(0, 0, 0, 0)):
     return new_image
 
 def check_input_image(input_image):
-    """Check if the input image is valid"""
     if input_image is None:
         raise gr.Error("No image uploaded!")
-    return input_image
+
 
 def remove_background(
     image: PIL.Image.Image,
@@ -111,7 +48,7 @@ def remove_background(
     do_remove = True
     if image.mode == "RGBA" and image.getextrema()[3][0] < 255:
         # explain why current do not rm bg
-        print("alpha channel not empty, skip remove background, using alpha channel as mask")
+        print("alhpa channl not enpty, skip remove background, using alpha channel as mask")
         background = Image.new("RGBA", image.size, (0, 0, 0, 0))
         image = Image.alpha_composite(background, image)
         do_remove = False
@@ -121,7 +58,7 @@ def remove_background(
     return image
 
 def do_resize_content(original_image: Image, scale_rate):
-    # resize image content while retaining the original image size
+    # resize image content wile retain the original image size
     if scale_rate != 1:
         # Calculate the new size after rescaling
         new_size = tuple(int(dim * scale_rate) for dim in original_image.size)
@@ -140,11 +77,6 @@ def add_background(image, bg_color=(255, 255, 255)):
     background = Image.new("RGBA", image.size, bg_color)
     return Image.alpha_composite(background, image)
 
-def add_random_background(image, color):
-    # Add a random background to the image
-    width, height = image.size
-    background = Image.new("RGBA", image.size, color)
-    return Image.alpha_composite(background, image)
 
 def preprocess_image(image, background_choice, foreground_ratio, backgroud_color):
     """
@@ -172,7 +104,53 @@ def gen_image(input_image, seed, scale, step):
     np_xyzs = np.concatenate(stage2_images, 1)
 
     glb_path = generate3d(model, np_imgs, np_xyzs, args.device)
-    return Image.fromarray(np_imgs), Image.fromarray(np_xyzs), glb_path
+    return Image.fromarray(np_imgs), Image.fromarray(np_xyzs), glb_path#, obj_path
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--stage1_config",
+    type=str,
+    default="configs/nf7_v3_SNR_rd_size_stroke.yaml",
+    help="config for stage1",
+)
+parser.add_argument(
+    "--stage2_config",
+    type=str,
+    default="configs/stage2-v2-snr.yaml",
+    help="config for stage2",
+)
+
+parser.add_argument("--device", type=str, default="cuda")
+args = parser.parse_args()
+
+crm_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="CRM.pth")
+specs = json.load(open("configs/specs_objaverse_total.json"))
+model = CRM(specs)
+model.load_state_dict(torch.load(crm_path, map_location="cpu"), strict=False)
+model = model.to(args.device)
+
+stage1_config = OmegaConf.load(args.stage1_config).config
+stage2_config = OmegaConf.load(args.stage2_config).config
+stage2_sampler_config = stage2_config.sampler
+stage1_sampler_config = stage1_config.sampler
+
+stage1_model_config = stage1_config.models
+stage2_model_config = stage2_config.models
+
+xyz_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="ccm-diffusion.pth")
+pixel_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="pixel-diffusion.pth")
+stage1_model_config.resume = pixel_path
+stage2_model_config.resume = xyz_path
+
+pipeline = TwoStagePipeline(
+    stage1_model_config,
+    stage2_model_config,
+    stage1_sampler_config,
+    stage2_sampler_config,
+    device=args.device,
+    dtype=torch.float32
+)
 
 _DESCRIPTION = '''
 * Our [official implementation](https://github.com/thu-ml/CRM) uses UV texture instead of vertex color. It has better texture than this online demo.
@@ -180,7 +158,6 @@ _DESCRIPTION = '''
 * If you find the output unsatisfying, try using different seeds:)
 '''
 
-# Create the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# CRM: Single Image to 3D Textured Mesh with Convolutional Reconstruction Model")
     gr.Markdown(_DESCRIPTION)
@@ -198,11 +175,13 @@ with gr.Blocks() as demo:
                 with gr.Column():
                     with gr.Row():
                         background_choice = gr.Radio([
-                            "Alpha as mask",
-                            "Auto Remove background"
-                        ], value="Auto Remove background",
-                        label="Background choice")
-                    back_ground_color = gr.ColorPicker(label="Background Color", value="#7F7F7F")
+                                "Alpha as mask",
+                                "Auto Remove background"
+                            ], value="Auto Remove background",
+                            label="backgroud choice")
+                        # do_remove_background = gr.Checkbox(label=, value=True)
+                        # force_remove = gr.Checkbox(label=, value=False)
+                    back_groud_color = gr.ColorPicker(label="Background Color", value="#7F7F7F", interactive=False)
                     foreground_ratio = gr.Slider(
                         label="Foreground Ratio",
                         minimum=0.5,
@@ -212,21 +191,21 @@ with gr.Blocks() as demo:
                     )
 
                 with gr.Column():
-                    seed = gr.Number(value=1234, label="Seed", precision=0)
-                    guidance_scale = gr.Number(value=5.5, minimum=3, maximum=10, label="Guidance scale")
-                    step = gr.Number(value=30, minimum=30, maximum=100, label="Sample steps", precision=0)
+                    seed = gr.Number(value=1234, label="seed", precision=0)
+                    guidance_scale = gr.Number(value=5.5, minimum=3, maximum=10, label="guidance_scale")
+                    step = gr.Number(value=30, minimum=30, maximum=100, label="sample steps", precision=0)
             text_button = gr.Button("Generate 3D shape")
-            # if os.path.exists("examples") and os.listdir("examples"):
-            #     gr.Examples(
-            #         examples=[os.path.join("examples", i) for i in os.listdir("examples") if i.lower().endswith(('.png', '.jpg', '.jpeg'))],
-            #         inputs=[image_input],
-            #         examples_per_page=20,
-            #     )
+            gr.Examples(
+                examples=[os.path.join("examples", i) for i in os.listdir("examples")],
+                inputs=[image_input],
+                examples_per_page = 20,
+            )
         with gr.Column():
             image_output = gr.Image(interactive=False, label="Output RGB image")
-            xyz_output = gr.Image(interactive=False, label="Output CCM image")
+            xyz_ouput = gr.Image(interactive=False, label="Output CCM image")
+
             output_model = gr.Model3D(
-                label="Output GLB",
+                label="Output OBJ",
                 interactive=False,
             )
             gr.Markdown("Note: Ensure that the input image is correctly pre-processed into a grey background, otherwise the results will be unpredictable.")
@@ -239,19 +218,19 @@ with gr.Blocks() as demo:
     ]
     outputs = [
         image_output,
-        xyz_output,
+        xyz_ouput,
         output_model,
+        # output_obj,
     ]
 
+
     text_button.click(fn=check_input_image, inputs=[image_input]).success(
         fn=preprocess_image,
-        inputs=[image_input, background_choice, foreground_ratio, back_ground_color],
+        inputs=[image_input, background_choice, foreground_ratio, back_groud_color],
         outputs=[processed_image],
     ).success(
         fn=gen_image,
         inputs=inputs,
         outputs=outputs,
     )
-
-if __name__ == "__main__":
-    demo.queue().launch()
+    demo.queue().launch()