Update app.py

app.py

@@ -76,55 +76,75 @@ def normalize_point_clouds(pcs, mode):
 def predict(Seed, ckpt):
     if Seed is None:
         Seed = 777
-    seed_all(int(Seed)) # Ensure Seed is an integer
-
-    # Ensure args is accessible, provide a default if it's missing or not a Namespace
-    # This is a defensive measure, as the error was about loading argparse.Namespace
-    if not hasattr(ckpt, 'args') or not hasattr(ckpt['args'], 'model'):
-        # This case should ideally not happen if the checkpoint is valid
-        # but if it does, we need a fallback or error.
-        # For now, let's assume 'args' and 'args.model' exist based on the error.
-        print("Warning: Checkpoint 'args' or 'args.model' not found. Assuming 'gaussian'.")
-        model_type = 'gaussian' 
-        latent_dim = ckpt.get('latent_dim', 128) # A common default
-        flexibility = ckpt.get('flexibility', 0.0) # A common default
+    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.
+    if 'args' in ckpt and hasattr(ckpt['args'], 'model'):
+        actual_args = ckpt['args']
+        print("Using 'args' found in checkpoint.")
     else:
-        model_type = ckpt['args'].model
-        latent_dim = ckpt['args'].latent_dim
-        flexibility = ckpt['args'].flexibility
-
-
-    if model_type == 'gaussian':
-        # Pass necessary args to the constructor
-        # We need to mock an args object if ckpt['args'] wasn't a full argparse.Namespace
-        # or if some attributes are missing.
-        mock_args = type('Args', (), {'latent_dim': latent_dim, 'hyper': getattr(ckpt.get('args', {}), 'hyper', None)})() # Add other required args
-        model = GaussianVAE(mock_args).to(device)
-    elif model_type == 'flow':
-        mock_args = type('Args', (), {
+        # Fallback to constructing a mock_args if 'args' is not as expected
+        # This part needs to be more robust and include all necessary defaults
+        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.
+        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_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
+        actual_args = type('Args', (), {
+            'model': model_type,
             'latent_dim': latent_dim,
-            'flow_depth': getattr(ckpt.get('args', {}), 'flow_depth', 10), # Example default
-            'flow_hidden_dim': getattr(ckpt.get('args', {}), 'flow_hidden_dim', 256), # Example default
-            'hyper': getattr(ckpt.get('args', {}), 'hyper', None)
+            '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 = FlowVAE(mock_args).to(device)
+    # --- 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':
+        model = FlowVAE(actual_args).to(device)
     else:
-        raise ValueError(f"Unknown model type: {model_type}")
+        raise ValueError(f"Unknown model type: {actual_args.model}")
 
     model.load_state_dict(ckpt['state_dict'])
-    model.eval() # Set model to evaluation mode
+    model.eval()
 
-    # Generate Point Clouds
     gen_pcs = []
     with torch.no_grad():
-        z = torch.randn([1, latent_dim]).to(device)
-        # The sample method might also depend on args from the checkpoint
-        num_points_to_generate = getattr(ckpt.get('args', {}), 'num_points', 2048) # Default to 2048 if not in args
-        x = model.sample(z, num_points_to_generate, flexibility=flexibility)
+        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())
-    
-    gen_pcs_tensor = torch.cat(gen_pcs, dim=0)[:1] # Ensure we take only one point cloud
-    gen_pcs_normalized = normalize_point_clouds(gen_pcs_tensor.clone(), mode="shape_bbox") # Use .clone() if normalize_point_clouds modifies inplace
+
+    gen_pcs_tensor = torch.cat(gen_pcs, dim=0)[:1]
+    gen_pcs_normalized = normalize_point_clouds(gen_pcs_tensor.clone(), mode="shape_bbox")
 
     return gen_pcs_normalized[0]