Update server.py

server.py

@@ -5,8 +5,7 @@ from flask_cors import CORS
 import io
 import os
 from PIL import Image
-from transformers import CLIPProcessor, CLIPModel
-import numpy as np
+from diffusers import StableDiffusionPipeline
 
 # Define the MIDM model
 class MIDM(nn.Module):
@@ -28,19 +27,18 @@ app = Flask(__name__, static_folder='static', template_folder='templates')
 CORS(app)
 
 # Load models once when the app starts to avoid reloading for each request
-processor = None
-clip_model = None
+stable_diff_pipe = None
 model = None
 
 def load_models():
-    global processor, clip_model, model
+    global stable_diff_pipe, model
     
-    # Load CLIP model and processor
-    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch14")
-    clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch14")
+    # Load Stable Diffusion model pipeline
+    stable_diff_pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4-original")
+    stable_diff_pipe.to("cuda" if torch.cuda.is_available() else "cpu")
     
     # Initialize MIDM model
-    input_dim = 10  # Using first 10 features as in your notebook
+    input_dim = 10  # Example dimension, adjust based on how you process the features
     hidden_dim = 64
     output_dim = 1
     model = MIDM(input_dim, hidden_dim, output_dim)
@@ -49,17 +47,19 @@ def load_models():
     # model.load_state_dict(torch.load('path/to/your/model.pth'))
     model.eval()
 
-# Function to get image features using CLIP
-def get_image_features(image):
+# Function to extract features from the image using Stable Diffusion
+def extract_image_features(image):
     """
-    Extracts image features using the CLIP model.
+    Extracts image features using the Stable Diffusion pipeline.
     """
-    # Preprocess the image and get features
-    inputs = processor(images=image, return_tensors="pt")
-    # Only use the image encoder to get the image features
+    # 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)
+    
+    # Generate the image embedding using the model
     with torch.no_grad():
-        image_features = clip_model.get_image_features(**inputs)
-    return image_features
+        generated_features = stable_diff_pipe.vae.encode(image_input).latent_dist.mean
+
+    return generated_features
 
 @app.route('/')
 def index():
@@ -68,7 +68,7 @@ def index():
 @app.route('/api/check-membership', methods=['POST'])
 def check_membership():
     # Ensure models are loaded
-    if processor is None or clip_model is None or model is None:
+    if stable_diff_pipe is None or model is None:
         load_models()
         
     if 'image' not in request.files:
@@ -79,12 +79,12 @@ def check_membership():
         file = request.files['image']
         image_bytes = file.read()
         image = Image.open(io.BytesIO(image_bytes))
-        
-        # Get image features using CLIP
-        image_features = get_image_features(image)
+
+        # Get image features using Stable Diffusion
+        image_features = extract_image_features(image)
         
         # Preprocess the features for MIDM model
-        processed_features = image_features.reshape(1, -1)[:, :10]  # Select first 10 features
+        processed_features = image_features.reshape(1, -1)[:, :10]  # Select first 10 features (example)
         
         # Perform inference
         with torch.no_grad():