Update gradio_app.py

gradio_app.py

@@ -47,12 +47,10 @@ def create_rgba_image(rgb_image: Image.Image, mask: np.ndarray = None) -> Image.
     """Create an RGBA image from RGB image and optional mask."""
     rgba_image = rgb_image.convert('RGBA')
     if mask is not None:
-        print("[debug] mask shape before alpha:", mask.shape)
         # Ensure mask is 2D before converting to alpha
         if len(mask.shape) > 2:
             mask = mask.squeeze()
         alpha = Image.fromarray((mask * 255).astype(np.uint8))
-        print("[debug] alpha size:", alpha.size)
         rgba_image.putalpha(alpha)
     return rgba_image
 
@@ -61,8 +59,7 @@ def create_batch(input_image: Image.Image) -> dict[str, Any]:
     # Resize and convert input image to numpy array
     resized_image = input_image.resize((COND_WIDTH, COND_HEIGHT))
     img_array = np.array(resized_image).astype(np.float32) / 255.0
-    print("[debug] img_array shape:", img_array.shape)
-    
+
     # Extract RGB and alpha channels
     if img_array.shape[-1] == 4:  # RGBA
         rgb = img_array[..., :3]
@@ -74,25 +71,18 @@ def create_batch(input_image: Image.Image) -> dict[str, Any]:
     # Convert to tensors while keeping channel-last format
     rgb = torch.from_numpy(rgb).float()  # [H, W, 3]
     mask = torch.from_numpy(mask).float()  # [H, W, 1]
-    print("[debug] rgb tensor shape:", rgb.shape)
-    print("[debug] mask tensor shape:", mask.shape)
-    
+
     # Create background blend (match channel-last format)
     bg_tensor = torch.tensor(BACKGROUND_COLOR).view(1, 1, 3)  # [1, 1, 3]
-    print("[debug] bg_tensor shape:", bg_tensor.shape)
-    
+ 
     # Blend RGB with background using mask (all in channel-last format)
     rgb_cond = torch.lerp(bg_tensor, rgb, mask)  # [H, W, 3]
-    print("[debug] rgb_cond shape after blend:", rgb_cond.shape)
-
+ 
     # Move channels to correct dimension and add batch dimension
     # Important: For SPAR3D image tokenizer, we need [B, H, W, C] format
     rgb_cond = rgb_cond.unsqueeze(0)  # [1, H, W, 3]
     mask = mask.unsqueeze(0)  # [1, H, W, 1]
     
-    print("[debug] rgb_cond final shape:", rgb_cond.shape)
-    print("[debug] mask final shape:", mask.shape)
-
     # Create the batch dictionary
     batch = {
         "rgb_cond": rgb_cond,  # [1, H, W, 3]
@@ -102,35 +92,20 @@ def create_batch(input_image: Image.Image) -> dict[str, Any]:
         "intrinsic_normed_cond": intrinsic_normed_cond.unsqueeze(0),  # [1, 3, 3]
     }
     
-    print("\nFinal batch shapes:")
     for k, v in batch.items():
         print(f"[debug] {k} final shape:", v.shape)
-    print("\nrgb_cond max:", batch["rgb_cond"].max())
-    print("rgb_cond min:", batch["rgb_cond"].min())
-    print("mask_cond unique values:", torch.unique(batch["mask_cond"]))
-        
+ 
     return batch
 
 def forward_model(batch, system, guidance_scale=3.0, seed=0, device="cuda"):
     """Process batch through model and generate point cloud."""
-    print("\n[debug] Starting forward_model")
-    print("[debug] Input rgb_cond shape:", batch["rgb_cond"].shape)
-    print("[debug] Input mask_cond shape:", batch["mask_cond"].shape)
-    
+
     batch_size = batch["rgb_cond"].shape[0]
     assert batch_size == 1, f"Expected batch size 1, got {batch_size}"
     
-    # Print value ranges for debugging
-    print("\nValue ranges:")
-    print("rgb_cond max:", batch["rgb_cond"].max())
-    print("rgb_cond min:", batch["rgb_cond"].min())
-    print("mask_cond unique values:", torch.unique(batch["mask_cond"]))
-    
     # Generate point cloud tokens
-    print("\n[debug] Generating point cloud tokens")
     try:
         cond_tokens = system.forward_pdiff_cond(batch)
-        print("[debug] cond_tokens shape:", cond_tokens.shape)
     except Exception as e:
         print("\n[ERROR] Failed in forward_pdiff_cond:")
         print(e)
@@ -141,7 +116,6 @@ def forward_model(batch, system, guidance_scale=3.0, seed=0, device="cuda"):
         raise
     
     # Sample points
-    print("\n[debug] Sampling points")
     sample_iter = system.sampler.sample_batch_progressive(
         batch_size,
         cond_tokens,
@@ -153,18 +127,14 @@ def forward_model(batch, system, guidance_scale=3.0, seed=0, device="cuda"):
     for x in sample_iter:
         samples = x["xstart"]
     
-    print("[debug] samples shape before permute:", samples.shape)
     pc_cond = samples.permute(0, 2, 1).float()
-    print("[debug] pc_cond shape after permute:", pc_cond.shape)
-    
+
     # Normalize point cloud
     pc_cond = spar3d_utils.normalize_pc_bbox(pc_cond)
-    print("[debug] pc_cond shape after normalize:", pc_cond.shape)
-    
+
     # Subsample to 512 points
     pc_cond = pc_cond[:, torch.randperm(pc_cond.shape[1])[:512]]
-    print("[debug] pc_cond final shape:", pc_cond.shape)
-    
+
     return pc_cond
 
 def generate_and_process_3d(prompt: str, seed: int = 42) -> tuple[str | None, Image.Image | None]:
@@ -180,7 +150,6 @@ def generate_and_process_3d(prompt: str, seed: int = 42) -> tuple[str | None, Im
         
         # Generate image using FLUX
         generator = torch.Generator(device=device).manual_seed(seed)
-        print("[debug] generating the image using Flux")
         generated_image = flux_pipe(
             prompt=prompt,
             width=width,
@@ -190,10 +159,8 @@ def generate_and_process_3d(prompt: str, seed: int = 42) -> tuple[str | None, Im
             guidance_scale=0.0
         ).images[0]
         
-        print("[debug] converting the image to rgb")
         rgb_image = generated_image.convert('RGB')
         
-        print("[debug] removing the background by calling bg_remover.process(rgb_image)")
         # bg_remover returns a PIL Image already, no need to convert
         no_bg_image = bg_remover.process(rgb_image)
         print(f"[debug] no_bg_image type: {type(no_bg_image)}, mode: {no_bg_image.mode}")
@@ -202,7 +169,6 @@ def generate_and_process_3d(prompt: str, seed: int = 42) -> tuple[str | None, Im
         rgba_image = no_bg_image.convert('RGBA')
         print(f"[debug] rgba_image mode: {rgba_image.mode}")
         
-        print("[debug] auto-cropping the rgba_image using spar3d_utils.foreground_crop(...)")
         processed_image = spar3d_utils.foreground_crop(
             rgba_image,
             crop_ratio=1.3,
@@ -215,7 +181,6 @@ def generate_and_process_3d(prompt: str, seed: int = 42) -> tuple[str | None, Im
         print(f"[debug] Alpha channel stats - min: {alpha.min()}, max: {alpha.max()}, unique: {np.unique(alpha)}")
 
         # Prepare batch for processing
-        print("[debug] preparing the batch by calling create_batch(processed_image)")
         batch = create_batch(processed_image)
         batch = {k: v.to(device) for k, v in batch.items()}
 
@@ -231,7 +196,6 @@ def generate_and_process_3d(prompt: str, seed: int = 42) -> tuple[str | None, Im
 
         # Generate mesh
         with torch.no_grad():
-            print("[debug] calling torch.autocast(....) to generate the mesh")
             with torch.autocast(device_type='cuda' if torch.cuda.is_available() else 'cpu', dtype=torch.bfloat16):
                 trimesh_mesh, _ = spar3d_model.generate_mesh(
                     batch,
@@ -243,20 +207,18 @@ def generate_and_process_3d(prompt: str, seed: int = 42) -> tuple[str | None, Im
                 trimesh_mesh = trimesh_mesh[0]
 
         # Export to GLB
-        print("[debug] creating tmp dir for the .glb output")
         temp_dir = tempfile.mkdtemp()
         output_path = os.path.join(temp_dir, 'output.glb')
         
-        print("[debug] calling trimesh_mesh.export(...) to export to .glb")
         trimesh_mesh.export(output_path, file_type="glb", include_normals=True)
         
-        return output_path, generated_image
+        return output_path
         
     except Exception as e:
         print(f"Error during generation: {str(e)}")
         import traceback
         traceback.print_exc()
-        return None, None
+        return None
         
 # Create Gradio interface
 demo = gr.Interface(
@@ -276,16 +238,12 @@ demo = gr.Interface(
     ],
     outputs=[
         gr.Model3D(
-            label="3D Model Preview",
+            label="Generated 3D model",
             clear_color=[0.0, 0.0, 0.0, 0.0],
-        ),
-        gr.Image(
-            label="Generated Image",
-            type="pil"
-        ),
+        )
     ],
-    title="Text to 3D Model Generator",
-    description="Enter a text prompt to generate an image that will be converted into a 3D model",
+    title="Text to 3D",
+    description="Enter a text prompt to generate an image that will be converted into a 3D model using Stable Point-Awaire 3D by Stability AI.",
 )
 
 if __name__ == "__main__":