Enhance app.py with improved error handling, compatibility fixes, and streamlined image processing functions. Update Gradio interface for better user experience, including clearer background choice options and refined button functionality.

app.py

@@ -1,4 +1,4 @@
-# Not ready to use yet
+# Fixed version with proper error handling and compatibility
 import spaces
 import argparse
 import numpy as np
@@ -36,7 +36,7 @@ def init_model():
         help="config for stage2",
     )
     parser.add_argument("--device", type=str, default="cuda")
-    args = parser.parse_args()
+    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")
@@ -73,6 +73,7 @@ def init_model():
 # Global variables to store model and pipeline
 model = None
 pipeline = None
+args = None
 
 @spaces.GPU
 def get_model():
@@ -84,7 +85,6 @@ def get_model():
 
 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
@@ -111,7 +111,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("alhpa channl not enpty, skip remove background, using alpha channel as mask")
+        print("alpha channel not empty, 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 +121,7 @@ def remove_background(
     return image
 
 def do_resize_content(original_image: Image, scale_rate):
-    # resize image content wile retain the original image size
+    # resize image content while retaining 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)
@@ -147,30 +147,35 @@ def add_random_background(image, color):
     return Image.alpha_composite(background, image)
 
 @spaces.GPU
-def preprocess_image(input_image, background_choice, foreground_ratio, back_groud_color):
+def preprocess_image(input_image, background_choice, foreground_ratio, back_ground_color):
     """Preprocess the input image"""
     try:
-        # Get model and pipeline when needed
-        model, pipeline = get_model()
+        # Check if image is provided
+        if input_image is None:
+            raise gr.Error("No image uploaded!")
+            
+        # Convert to PIL Image if needed
+        if isinstance(input_image, np.ndarray):
+            input_image = Image.fromarray(input_image)
         
-        # Convert to numpy array
-        np_image = np.array(input_image)
+        # Ensure RGBA mode
+        if input_image.mode != "RGBA":
+            input_image = input_image.convert("RGBA")
         
         # Process background
-        if background_choice == "Remove Background":
-            np_image = rembg.remove(np_image, session=rembg_session)
+        if background_choice == "Auto Remove background":
+            input_image = remove_background(input_image, rembg_session)
         elif background_choice == "Custom Background":
-            np_image = add_random_background(np_image, back_groud_color)
+            input_image = add_random_background(input_image, back_ground_color)
             
         # Resize content if needed
         if foreground_ratio != 1.0:
-            np_image = do_resize_content(Image.fromarray(np_image), foreground_ratio)
-            np_image = np.array(np_image)
+            input_image = do_resize_content(input_image, foreground_ratio)
             
-        return Image.fromarray(np_image)
+        return input_image
     except Exception as e:
         print(f"Error in preprocess_image: {str(e)}")
-        raise e
+        raise gr.Error(f"Preprocessing failed: {str(e)}")
 
 @spaces.GPU
 def gen_image(processed_image, seed, scale, step):
@@ -179,26 +184,51 @@ def gen_image(processed_image, seed, scale, step):
         # Get model and pipeline when needed
         model, pipeline = get_model()
         
+        # Check if image is provided
+        if processed_image is None:
+            raise gr.Error("No processed image provided!")
+        
         # Convert to numpy array
-        np_image = np.array(processed_image)
+        if isinstance(processed_image, Image.Image):
+            np_image = np.array(processed_image)
+        else:
+            np_image = processed_image
         
         # Set random seed
-        torch.manual_seed(seed)
-        np.random.seed(seed)
+        torch.manual_seed(int(seed))
+        np.random.seed(int(seed))
         
         # Generate images
         np_imgs, np_xyzs = pipeline.generate(
             np_image,
-            guidance_scale=scale,
-            num_inference_steps=step
+            guidance_scale=float(scale),
+            num_inference_steps=int(step)
         )
         
         # Generate 3D model
         glb_path = generate3d(model, np_imgs, np_xyzs, args.device)
+        
         return Image.fromarray(np_imgs), Image.fromarray(np_xyzs), glb_path
     except Exception as e:
         print(f"Error in gen_image: {str(e)}")
-        raise e
+        raise gr.Error(f"Generation failed: {str(e)}")
+
+def process_and_generate(image, bg_choice, fg_ratio, bg_color, seed_val, guidance, steps):
+    """Combined function to process image and generate 3D model"""
+    try:
+        if image is None:
+            raise gr.Error("No image uploaded!")
+        
+        # Preprocess the image
+        processed = preprocess_image(image, bg_choice, fg_ratio, bg_color)
+        
+        # Generate 3D model
+        rgb_img, ccm_img, glb_file = gen_image(processed, seed_val, guidance, steps)
+        
+        return processed, rgb_img, ccm_img, glb_file
+    except Exception as e:
+        print(f"Error in process_and_generate: {str(e)}")
+        raise gr.Error(f"Process failed: {str(e)}")
 
 _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.
@@ -206,7 +236,8 @@ _DESCRIPTION = '''
 * If you find the output unsatisfying, try using different seeds:)
 '''
 
-with gr.Blocks() as demo:
+# Create the Gradio interface
+with gr.Blocks(title="CRM: Single Image to 3D") as demo:
     gr.Markdown("# CRM: Single Image to 3D Textured Mesh with Convolutional Reconstruction Model")
     gr.Markdown(_DESCRIPTION)
     
@@ -215,21 +246,23 @@ with gr.Blocks() as demo:
             with gr.Row():
                 image_input = gr.Image(
                     label="Image input",
-                    type="pil"
+                    type="pil",
+                    height=300
                 )
                 processed_image = gr.Image(
                     label="Processed Image",
-                    type="pil"
+                    type="pil",
+                    height=300
                 )
             
             with gr.Row():
                 with gr.Column():
                     background_choice = gr.Radio(
-                        choices=["Alpha as mask", "Auto Remove background"],
+                        choices=["Alpha as mask", "Auto Remove background", "Custom Background"],
                         value="Auto Remove background",
                         label="Background choice"
                     )
-                    back_groud_color = gr.ColorPicker(
+                    back_ground_color = gr.ColorPicker(
                         label="Background Color",
                         value="#7F7F7F"
                     )
@@ -261,12 +294,14 @@ with gr.Blocks() as demo:
                         precision=0
                     )
             
-            text_button = gr.Button("Generate 3D shape")
+            text_button = gr.Button("Generate 3D shape", variant="primary")
             
-            gr.Examples(
-                examples=[[os.path.join("examples", i)] for i in os.listdir("examples")],
-                inputs=[image_input]
-            )
+            # Only add examples if the directory exists
+            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]
+                )
         
         with gr.Column():
             image_output = gr.Image(label="Output RGB image")
@@ -274,24 +309,20 @@ with gr.Blocks() as demo:
             output_model = gr.Model3D(label="Output GLB")
             gr.Markdown("Note: Ensure that the input image is correctly pre-processed into a grey background, otherwise the results will be unpredictable.")
 
-    def process_and_generate(image, bg_choice, fg_ratio, bg_color, seed_val, guidance, steps):
-        if image is None:
-            raise gr.Error("No image uploaded!")
-        processed = preprocess_image(image, bg_choice, fg_ratio, bg_color)
-        return gen_image(processed, seed_val, guidance, steps)
-
+    # Connect the button to the processing function
     text_button.click(
         fn=process_and_generate,
         inputs=[
             image_input,
             background_choice,
             foreground_ratio,
-            back_groud_color,
+            back_ground_color,
             seed,
             guidance_scale,
             step
         ],
         outputs=[
+            processed_image,
             image_output,
             xyz_output,
             output_model
@@ -300,4 +331,10 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    # Launch with proper settings for HF Spaces
+    demo.queue().launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,  # Don't need share=True on HF Spaces
+        show_api=False  # Disable API docs to avoid the error
+    )