Update app.py

app.py

@@ -79,54 +79,58 @@ def predict(Seed, ckpt):
     seed_all(int(Seed))
 
     # --- MODIFICATION START ---
-    # Try to get the original args from the checkpoint first
-    # The key might be 'args', 'config', or something similar.
-    # We need to inspect the actual keys of a loaded ckpt if this doesn't work.
+    actual_args = None
     if 'args' in ckpt and hasattr(ckpt['args'], 'model'):
         actual_args = ckpt['args']
         print("Using 'args' found in checkpoint.")
     else:
-        # Fallback to constructing a mock_args if 'args' is not as expected
-        # This part needs to be more robust and include all necessary defaults
+        # This fallback should ideally not be hit if 'args' is usually present
         print("Warning: 'args' not found or 'args.model' missing in checkpoint. Constructing mock_args.")
-        
-        # Defaults - these might need to be adjusted based on the original training scripts
-        # or by inspecting a correctly loaded checkpoint from the original repo.
+        # Define all necessary defaults if we have to construct from scratch
+        default_model_type = 'gaussian'
         default_latent_dim = 128
-        default_hyper = None # Or some sensible default if PointwiseNet/etc. need it
-        default_residual = True # Common default for PointwiseNet, but needs verification
+        default_hyper = None
+        default_residual = True
         default_flow_depth = 10
         default_flow_hidden_dim = 256
-        default_model_type = 'gaussian' # Default if not found
-        default_num_points = 2048
-        default_flexibility = 0.0
-
-        # Try to get values from ckpt if they exist at the top level
-        # (some checkpoints might store them flatly instead of under an 'args' key)
-        model_type = ckpt.get('model', default_model_type) # Check if 'model' key exists directly
-        latent_dim = ckpt.get('latent_dim', default_latent_dim)
-        hyper = ckpt.get('hyper', default_hyper)
-        residual = ckpt.get('residual', default_residual)
-        flow_depth = ckpt.get('flow_depth', default_flow_depth)
-        flow_hidden_dim = ckpt.get('flow_hidden_dim', default_flow_hidden_dim)
-        num_points_to_generate = ckpt.get('num_points', default_num_points)
-        flexibility = ckpt.get('flexibility', default_flexibility)
-
-        # Create the mock_args object
+        default_num_points = 2048 # Default for sampling
+        default_flexibility = 0.0  # Default for sampling
+
         actual_args = type('Args', (), {
-            'model': model_type,
-            'latent_dim': latent_dim,
-            'hyper': hyper,
-            'residual': residual, # Added residual
-            'flow_depth': flow_depth,
-            'flow_hidden_dim': flow_hidden_dim,
-            'num_points': num_points_to_generate,
-            'flexibility': flexibility
-            # Add any other attributes that models might expect from 'args'
+            'model': ckpt.get('model', default_model_type),
+            'latent_dim': ckpt.get('latent_dim', default_latent_dim),
+            'hyper': ckpt.get('hyper', default_hyper),
+            'residual': ckpt.get('residual', default_residual),
+            'flow_depth': ckpt.get('flow_depth', default_flow_depth),
+            'flow_hidden_dim': ckpt.get('flow_hidden_dim', default_flow_hidden_dim),
+            'num_points': ckpt.get('num_points', default_num_points), # Try to get from ckpt top-level too
+            'flexibility': ckpt.get('flexibility', default_flexibility) # Try to get from ckpt top-level too
         })()
+
+    # Ensure essential attributes for sampling exist on actual_args, even if 'args' was found
+    # These are parameters for the .sample() method, not necessarily model construction.
+    # The original training args might not have included these if they were fixed in the sampling script.
+    
+    # Default values for sampling parameters if not present in actual_args
+    default_num_points_sampling = 2048
+    default_flexibility_sampling = 0.0
+
+    if not hasattr(actual_args, 'num_points'):
+        print(f"Attribute 'num_points' not found in actual_args. Setting default: {default_num_points_sampling}")
+        setattr(actual_args, 'num_points', default_num_points_sampling)
+    
+    if not hasattr(actual_args, 'flexibility'):
+        print(f"Attribute 'flexibility' not found in actual_args. Setting default: {default_flexibility_sampling}")
+        setattr(actual_args, 'flexibility', default_flexibility_sampling)
+        
+    # Also ensure 'residual' is present if it's a Gaussian model, as it was an issue before
+    # This is more for model construction, but good to double-check if the 'args' from ckpt might be incomplete
+    if actual_args.model == 'gaussian' and not hasattr(actual_args, 'residual'):
+        print(f"Attribute 'residual' not found in actual_args for Gaussian model. Setting default: True")
+        setattr(actual_args, 'residual', True) # Default for GaussianVAE
+
     # --- MODIFICATION END ---
 
-    # Now use actual_args to instantiate models
     if actual_args.model == 'gaussian':
         model = GaussianVAE(actual_args).to(device)
     elif actual_args.model == 'flow':
@@ -139,6 +143,7 @@ def predict(Seed, ckpt):
 
     gen_pcs = []
     with torch.no_grad():
+        # Use the (potentially now augmented) actual_args for sampling
         z = torch.randn([1, actual_args.latent_dim]).to(device)
         x = model.sample(z, actual_args.num_points, flexibility=actual_args.flexibility)
         gen_pcs.append(x.detach().cpu())
@@ -147,7 +152,6 @@ def predict(Seed, ckpt):
     gen_pcs_normalized = normalize_point_clouds(gen_pcs_tensor.clone(), mode="shape_bbox")
 
     return gen_pcs_normalized[0]
-
 def generate(seed, value):
     if value == "Airplane":
         ckpt = ckpt_airplane