Update server.py

server.py

@@ -33,11 +33,11 @@ CORS(app)
 stable_diff_pipe = None
 model = None
 
-def load_models():
+def load_models(model_name="CompVis/stable-diffusion-v1-4"):
     global stable_diff_pipe, model
     
     # Load Stable Diffusion model pipeline
-    stable_diff_pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
+    stable_diff_pipe = StableDiffusionPipeline.from_pretrained(model_name)
     stable_diff_pipe.to("cuda" if torch.cuda.is_available() else "cpu")
     
     # Initialize MIDM model
@@ -66,9 +66,15 @@ def index():
 
 @app.route('/api/check-membership', methods=['POST'])
 def check_membership():
-    # Ensure models are loaded
+    # 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()
+        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