Upload: load.py

single/load.py

@@ -71,20 +71,24 @@ class EnsembleModel(torch.nn.Module):
             uncertainty = torch.zeros_like(output_probs)
         
         return output_probs, uncertainty  # Both (B, 1, H, W)
-    
+
 def compiled_model(
     path: pathlib.Path, 
     stac_item: pystac.Item,
     mode: Literal["min", "mean", "median", "max"] = "max",
     *args, **kwargs
-) -> EnsembleModel:
+):
     """
-    Loads the ensemble dynamically using the 'stac_item' as the source of truth.
+    Loads model(s) dynamically based on STAC metadata.
+    
+    - If single .pt2 → returns single model
+    - If multiple .pt2 → returns EnsembleModel
     
     Args:
-        mode: Aggregation mode - 'min', 'mean', 'median', or 'max' (default: 'max')
+        mode: Aggregation mode for ensembles (ignored for single models)
     
-    Returns an EnsembleModel that outputs (probabilities, uncertainty) tuple.
+    Returns:
+        Single model or EnsembleModel
     """
     model_paths = []
     for asset_key, asset in stac_item.assets.items():
@@ -93,10 +97,16 @@ def compiled_model(
     
     if not model_paths:
         raise ValueError("No .pt2 files found in STAC item assets.")
+    
     model_paths.sort()
-    models = [torch.export.load(p).module() for p in model_paths]
     
-    return EnsembleModel(*models, mode=mode)
+    if len(model_paths) == 1:
+        # Single model
+        return torch.export.load(model_paths[0]).module()
+    else:
+        # Ensemble model
+        models = [torch.export.load(p).module() for p in model_paths]
+        return EnsembleModel(*models, mode=mode)
 
 def define_iteration(dimension: tuple, chunk_size: int, overlap: int = 0):
     """
@@ -172,8 +182,7 @@ def predict_large(
     except (NotImplementedError, AttributeError):
         # Exported model (.pt2) or EnsembleModel
         model = model.to(device)
-
-    # Test if model returns tuple (ensemble) or single output
+        
     test_input = torch.zeros(1, bands, chunk_size, chunk_size).to(device)
     with torch.no_grad():
         test_output = model(test_input)
@@ -300,8 +309,7 @@ def predict_large(
                 offset_y : offset_y + length_y,
                 offset_x : offset_x + length_x
             ] = to_write_uncertainty
-
-    # Return based on model type
+            
     if is_ensemble:
         return output_probs, output_uncertainty
     else: