Update app.py

app.py

@@ -14,20 +14,23 @@ from PIL import Image
 import sf3d.utils as sf3d_utils
 from sf3d.system import SF3D
 
+# Initialize the rembg session
 rembg_session = rembg.new_session()
 
+# Constants
 COND_WIDTH = 512
 COND_HEIGHT = 512
 COND_DISTANCE = 1.6
 COND_FOVY_DEG = 40
 BACKGROUND_COLOR = [0.5, 0.5, 0.5]
 
-# Cached. Doesn't change
+# Cached camera parameters
 c2w_cond = sf3d_utils.default_cond_c2w(COND_DISTANCE)
 intrinsic, intrinsic_normed_cond = sf3d_utils.create_intrinsic_from_fov_deg(
     COND_FOVY_DEG, COND_HEIGHT, COND_WIDTH
 )
 
+# Load the model
 model = SF3D.from_pretrained(
     "stabilityai/stable-fast-3d",
     config_name="config.yaml",
@@ -35,11 +38,12 @@ model = SF3D.from_pretrained(
 )
 model.eval().cuda()
 
+# Load example files
 example_files = [
     os.path.join("demo_files/examples", f) for f in os.listdir("demo_files/examples")
 ]
 
-
+# Define functions
 def run_model(input_image):
     start = time.time()
     with torch.no_grad():
@@ -49,16 +53,13 @@ def run_model(input_image):
             trimesh_mesh, _glob_dict = model.generate_mesh(model_batch, 1024)
             trimesh_mesh = trimesh_mesh[0]
 
-    # Create new tmp file
+    # Create new temporary file
     tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".glb")
-
     trimesh_mesh.export(tmp_file.name, file_type="glb", include_normals=True)
 
     print("Generation took:", time.time() - start, "s")
-
     return tmp_file.name
 
-
 def create_batch(input_image: Image) -> dict[str, Any]:
     img_cond = (
         torch.from_numpy(
@@ -80,11 +81,10 @@ def create_batch(input_image: Image) -> dict[str, Any]:
         "intrinsic_cond": intrinsic.unsqueeze(0),
         "intrinsic_normed_cond": intrinsic_normed_cond.unsqueeze(0),
     }
-    # Add batch dim
+    # Add batch dimension
     batched = {k: v.unsqueeze(0) for k, v in batch_elem.items()}
     return batched
 
-
 @lru_cache
 def checkerboard(squares: int, size: int, min_value: float = 0.5):
     base = np.zeros((squares, squares)) + min_value
@@ -98,11 +98,9 @@ def checkerboard(squares: int, size: int, min_value: float = 0.5):
         .repeat(3, axis=-1)
     )
 
-
 def remove_background(input_image: Image) -> Image:
     return rembg.remove(input_image, session=rembg_session)
 
-
 def resize_foreground(
     image: Image,
     ratio: float,
@@ -126,26 +124,25 @@ def resize_foreground(
         fg,
         ((ph0, ph1), (pw0, pw1), (0, 0)),
         mode="constant",
-        constant_values=((0, 0), (0, 0), (0, 0)),
+        constant_values=0,
     )
 
     # Compute padding according to the ratio
     new_size = int(new_image.shape[0] / ratio)
-    # Pad to size, double side
+    # Pad to new size
     ph0, pw0 = (new_size - size) // 2, (new_size - size) // 2
     ph1, pw1 = new_size - size - ph0, new_size - size - pw0
     new_image = np.pad(
         new_image,
         ((ph0, ph1), (pw0, pw1), (0, 0)),
         mode="constant",
-        constant_values=((0, 0), (0, 0), (0, 0)),
+        constant_values=0,
     )
     new_image = Image.fromarray(new_image, mode="RGBA").resize(
         (COND_WIDTH, COND_HEIGHT)
     )
     return new_image
 
-
 def square_crop(input_image: Image) -> Image:
     # Perform a center square crop
     min_size = min(input_image.size)
@@ -157,7 +154,6 @@ def square_crop(input_image: Image) -> Image:
         (COND_WIDTH, COND_HEIGHT)
     )
 
-
 def show_mask_img(input_image: Image) -> Image:
     img_numpy = np.array(input_image)
     alpha = img_numpy[:, :, 3] / 255.0
@@ -165,11 +161,9 @@ def show_mask_img(input_image: Image) -> Image:
     new_img = img_numpy[..., :3] * alpha[:, :, None] + chkb * (1 - alpha[:, :, None])
     return Image.fromarray(new_img.astype(np.uint8), mode="RGB")
 
-
 def run_button(run_btn, input_image, background_state, foreground_ratio):
     if run_btn == "Run":
         glb_file: str = run_model(background_state)
-
         return (
             gr.update(),
             gr.update(),
@@ -180,10 +174,8 @@ def run_button(run_btn, input_image, background_state, foreground_ratio):
         )
     elif run_btn == "Remove Background":
         rem_removed = remove_background(input_image)
-
         sqr_crop = square_crop(rem_removed)
         fr_res = resize_foreground(sqr_crop, foreground_ratio)
-
         return (
             gr.update(value="Run", visible=True),
             sqr_crop,
@@ -193,7 +185,6 @@ def run_button(run_btn, input_image, background_state, foreground_ratio):
             gr.update(visible=False),
         )
 
-
 def requires_bg_remove(image, fr):
     if image is None:
         return (
@@ -228,7 +219,6 @@ def requires_bg_remove(image, fr):
         gr.update(visible=False),
     )
 
-
 def update_foreground_ratio(img_proc, fr):
     foreground_res = resize_foreground(img_proc, fr)
     return (
@@ -236,38 +226,27 @@ def update_foreground_ratio(img_proc, fr):
         gr.update(value=show_mask_img(foreground_res)),
     )
 
-# Generate color shades for the primary hue
-from gradio.themes.colors import Color
-
-primary_hue_color = Color(
-    name="custom",
-    c50="#e7e7e8",
-    c100="#c2c2c4",
-    c200="#9d9da0",
-    c300="#78787b",
-    c400="#525357",
-    c500="#2d2e33",
-    c600="#191a1e",  # Base color
-    c700="#141517",
-    c800="#0f1012",
-    c900="#0a0a0c",
-    c950="#050506",
-)
-
-# Define the custom theme
+# Define custom theme
 class CustomTheme(gr.themes.Base):
-    primary_hue = primary_hue_color
-    background_fill_primary = "#191a1e"
-    background_fill_secondary = "#191a1e"
-    background_fill_tertiary = "#191a1e"
-    text_color_primary = "#FFFFFF"
-    text_color_secondary = "#FFFFFF"
-    text_color_tertiary = "#FFFFFF"
-    input_background_fill = "#191a1e"
-    input_text_color = "#FFFFFF"
+    def __init__(self):
+        super().__init__()
+        self.primary_hue = "#191a1e"
+        self.background_fill_primary = "#191a1e"
+        self.background_fill_secondary = "#191a1e"
+        self.background_fill_tertiary = "#191a1e"
+        self.text_color_primary = "#FFFFFF"
+        self.text_color_secondary = "#FFFFFF"
+        self.text_color_tertiary = "#FFFFFF"
+        self.input_background_fill = "#191a1e"
+        self.input_text_color = "#FFFFFF"
+        self.font = (
+            "Poppins",
+            "https://fonts.googleapis.com/css2?family=Poppins:wght@400;500;700&display=swap",
+        )
 
+# Custom CSS
 css = """
-/* Apply background color to the entire page */
+/* Set background color for the entire page */
 body {
     background-color: #191a1e !important;
     margin: 0;
@@ -275,43 +254,46 @@ body {
 }
 
 /* Apply fonts */
-@import url('https://fonts.googleapis.com/css2?family=Poppins:wght@400;500;700&display=swap');
 body, input, button, textarea, select, .gr-button {
     font-family: 'Poppins', sans-serif;
     color: #FFFFFF;
 }
 
-/* Header settings */
+/* Header styles */
 h1, h2, h3, h4, h5, h6 {
     font-family: 'Poppins', sans-serif;
     font-weight: 700;
     color: #FFFFFF;
 }
 
-/* Button colors */
-.gr-button {
+/* Button styles */
+.generate-button {
     background-color: #5271FF !important;
     color: #FFFFFF !important;
     border: none;
     font-weight: bold;
 }
 
-/* Container colors */
+/* Container styles */
 .gr-block {
     background-color: #1c1c24 !important;
     border: 1px solid #5271FF !important;
 }
 """
 
+# Build the Gradio interface
 with gr.Blocks(theme=CustomTheme(), css=css) as demo:
     img_proc_state = gr.State()
     background_remove_state = gr.State()
-  
+
     with gr.Row(variant="panel"):
         with gr.Column():
             with gr.Row():
                 input_img = gr.Image(
-                    type="pil", label="Input Image", source="upload", image_mode="RGBA"
+                    type="pil",
+                    label="Input Image",
+                    sources="upload",
+                    image_mode="RGBA",
                 )
                 preview_removal = gr.Image(
                     label="Preview Background Removal",
@@ -335,7 +317,12 @@ with gr.Blocks(theme=CustomTheme(), css=css) as demo:
                 outputs=[background_remove_state, preview_removal],
             )
 
-            run_btn = gr.Button("Run", variant="primary", visible=False)
+            run_btn = gr.Button(
+                "Run",
+                variant="primary",
+                visible=False,
+                elem_classes="generate-button",
+            )
 
         with gr.Column():
             output_3d = LitModel3D(
@@ -347,10 +334,13 @@ with gr.Blocks(theme=CustomTheme(), css=css) as demo:
                 scale=1.0,
             )
             with gr.Column(visible=False, scale=1.0) as hdr_row:
-                gr.Markdown("""## HDR Environment Map
+                gr.Markdown(
+                    """
+                    ## HDR Environment Map
 
-Select an HDR environment map to light the 3D model. You can also upload your own HDR environment maps.
-""")
+                    Select an HDR environment map to light the 3D model. You can also upload your own HDR environment maps.
+                    """
+                )
 
                 with gr.Row():
                     hdr_illumination_file = gr.File(
@@ -407,4 +397,5 @@ Select an HDR environment map to light the 3D model. You can also upload your ow
         ],
     )
 
+# Launch the interface
 demo.launch()