Update server.py

server.py

@@ -41,20 +41,16 @@ def load_models():
     stable_diff_pipe.to("cuda" if torch.cuda.is_available() else "cpu")
     
     # Initialize MIDM model
-    input_dim = 10  # Example dimension, adjust based on how you process the features
+    input_dim = 10  
     hidden_dim = 64
     output_dim = 1
     model = MIDM(input_dim, hidden_dim, output_dim)
     
-    # For a real application, you would load your trained weights here
-    # model.load_state_dict(torch.load('path/to/your/model.pth'))
     model.eval()
 
 # Function to extract features from the image using Stable Diffusion
 def extract_image_features(image):
-    """
-    Extracts image features using the Stable Diffusion pipeline.
-    """
+    #Extracts image features using the Stable Diffusion pipeline.
     # Preprocess the image and get the feature vector
     image_input = stable_diff_pipe.feature_extractor(image, return_tensors="pt").pixel_values.to(stable_diff_pipe.device)
     
@@ -87,7 +83,7 @@ def check_membership():
         image_features = extract_image_features(image)
         
         # Preprocess the features for MIDM model
-        processed_features = image_features.reshape(1, -1)[:, :10]  # Select first 10 features (example)
+        processed_features = image_features.reshape(1, -1)[:, :10]  # Select first 10 features
         
         # Perform inference
         with torch.no_grad():