Update server.py

server.py

@@ -33,15 +33,11 @@ CORS(app)
 stable_diff_pipe = None
 model = None
 
-def load_models(model_name="CompVis/stable-diffusion-v1-4"):
+def load_models(model_name="rupeshs/LCM-runwayml-stable-diffusion-v1-5"):
     global stable_diff_pipe, model
     
     # Load Stable Diffusion model pipeline
-    stable_diff_pipe = StableDiffusionPipeline.from_pretrained(
-        model_name,
-        safety_checker=None,  # Disable safety checker for feature extraction
-        requires_safety_checker=False  # Explicitly indicate we don't need the safety checker
-    )
+    stable_diff_pipe = StableDiffusionPipeline.from_pretrained(model_name)
     stable_diff_pipe.to("cuda" if torch.cuda.is_available() else "cpu")
     
     # Initialize MIDM model
@@ -70,20 +66,20 @@ def index():
 
 @app.route('/api/check-membership', methods=['POST'])
 def check_membership():
-    # Get the model name from the request
-    model_name = request.form.get('model', 'CompVis/stable-diffusion-v1-4')
-    
-    # Ensure models are loaded with the selected model
-    if stable_diff_pipe is None or model is None:
-        load_models(model_name)
-    elif stable_diff_pipe.name_or_path != model_name:
-        # Reload the model if a different one is selected
-        load_models(model_name)
-        
-    if 'image' not in request.files:
-        return jsonify({'error': 'No image found in request'}), 400
-    
     try:
+        # Get the model name from the request
+        model_name = request.form.get('model', 'rupeshs/LCM-runwayml-stable-diffusion-v1-5')
+        
+        # Ensure models are loaded with the selected model
+        if stable_diff_pipe is None or model is None:
+            load_models(model_name)
+        elif stable_diff_pipe.name_or_path != model_name:
+            # Reload the model if a different one is selected
+            load_models(model_name)
+            
+        if 'image' not in request.files:
+            return jsonify({'error': 'No image found in request'}), 400
+        
         # Get the image from the request
         file = request.files['image']
         image_bytes = file.read()
@@ -109,6 +105,7 @@ def check_membership():
         })
         
     except Exception as e:
+        print(f"Error processing request: {str(e)}")  # Add logging
         return jsonify({'error': str(e)}), 500
 
 if __name__ == '__main__':