Update app.py

app.py

@@ -5,6 +5,8 @@ import torch
 import numpy as np
 import io
 import base64
+import threading
+import time
 
 app = Flask(__name__)
 
@@ -12,6 +14,9 @@ app = Flask(__name__)
 processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
 model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
 
+# Global variable for caching results
+cache = {}
+
 # Function to process image and generate mask
 def process_image(image, prompt):
     inputs = processor(
@@ -47,6 +52,11 @@ def get_masks(prompts, img, threshold):
 
 # Function to extract image using positive and negative prompts
 def extract_image(pos_prompts, neg_prompts, img, threshold):
+    cache_key = (pos_prompts, neg_prompts, threshold)
+    
+    if cache_key in cache:
+        return cache[cache_key]
+
     positive_masks = get_masks(pos_prompts, img, 0.5)
     negative_masks = get_masks(neg_prompts, img, 0.5)
 
@@ -55,15 +65,16 @@ def extract_image(pos_prompts, neg_prompts, img, threshold):
     final_mask = pos_mask & ~neg_mask
 
     final_mask = Image.fromarray(final_mask.astype(np.uint8) * 255, "L")
-    output_image = Image.new("RGBA", img.size, (0, 0, 0, 0))
-    output_image.paste(img, mask=final_mask)
 
     # Convert final mask to base64
     buffered = io.BytesIO()
     final_mask.save(buffered, format="PNG")
     final_mask_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
 
-    return final_mask_base64
+    # Cache the result
+    cache[cache_key] = {'final_mask_base64': final_mask_base64}
+
+    return cache[cache_key]
 
 @app.route('/')
 def hello_world():
@@ -84,10 +95,21 @@ def process_request():
     threshold = float(data.get('threshold', 0.4))
 
     # Perform image segmentation
-    final_mask_base64 = extract_image(pos_prompts, neg_prompts, img, threshold)
+    result = extract_image(pos_prompts, neg_prompts, img, threshold)
 
-    return jsonify({'final_mask_base64': final_mask_base64})
+    return jsonify(result)
+
+# Keep the server alive using a periodic task
+def keep_alive():
+    while True:
+        time.sleep(300)  # 5 minutes
+        requests.get('http://127.0.0.1:7860/')  # Send a request to keep the server alive
 
 if __name__ == '__main__':
     print("Server starting. Verify it is running by visiting http://0.0.0.0:7860/")
+
+    # Start the keep-alive thread
+    keep_alive_thread = threading.Thread(target=keep_alive)
+    keep_alive_thread.start()
+
     app.run(host='0.0.0.0', port=7860, debug=True)