data

starvector/serve/vllm_api_gradio/gradio_web_server.py

@@ -8,6 +8,13 @@ import requests
 from starvector.serve.conversation import default_conversation
 from starvector.serve.constants import LOGDIR, CLIP_QUERY_LENGTH
 from starvector.serve.util import (build_logger, server_error_msg)
+import boto3
+import hashlib
+from io import BytesIO
+
+# Configure your S3 bucket and key prefix (you can also use environment variables)
+S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME", "your-bucket-name")
+S3_KEY_PREFIX = os.getenv("S3_KEY_PREFIX", "uploads")
 
 logger = build_logger("gradio_web_server", "gradio_web_server.log")
 headers = {"User-Agent": "StarVector Client"}
@@ -20,6 +27,11 @@ priority = {
     "starvector-1b": "aaaaaaa",
 }
 
+
+# Configure your S3 bucket and key prefix (you can also use environment variables)
+S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME", "your-bucket-name")
+S3_KEY_PREFIX = os.getenv("S3_KEY_PREFIX", "starvector_uploads")
+
 def get_conv_log_filename():
     t = datetime.datetime.now()
     name = os.path.join(LOGDIR, f"{t.year}-{t.month:02d}-{t.day:02d}-conv.json")
@@ -108,6 +120,22 @@ def clear_history(request: gr.Request):
     state = default_conversation.copy()
     return (state, None, None) + (disable_btn,) * 7
 
+def upload_image_to_s3(image, bucket_name=S3_BUCKET_NAME, key_prefix=S3_KEY_PREFIX):
+    """
+    Upload a PIL Image to S3.
+    """
+    s3 = boto3.client("s3")
+    buffer = BytesIO()
+    # Save image to buffer; adjust format as needed (e.g., "PNG" or "JPEG")
+    image.save(buffer, format="JPEG")
+    buffer.seek(0)
+    # Create a unique S3 key using a timestamp
+    timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+    img_md5 = str(hashlib.md5(buffer.getvalue()).hexdigest())
+    s3_key = f"{key_prefix}/image_{img_md5}.jpeg"
+    s3.upload_fileobj(buffer, bucket_name, s3_key)
+    return s3_key
+
 def send_data(state, image, image_process_mode, text_caption, task, request: gr.Request):
     logger.info(f"send_data. ip: {request.client.host}.")
     if task == 'Image2SVG':
@@ -115,11 +143,15 @@ def send_data(state, image, image_process_mode, text_caption, task, request: gr.
             state.skip_next = True
             return (state, None, None, image) + (no_change_btn,) * 7
 
+        # Upload the image to S3
+        s3_key = upload_image_to_s3(image)
+        logger.info(f"Image uploaded to S3 at key: {s3_key}")
+
         # Reset the conversation state when a new image is uploaded
         state = default_conversation.copy()
         
-        if image is not None:
-            image_message = (image, image_process_mode)
+        # Prepare image message (you can also include the S3 URL or key if needed)
+        image_message = (image, image_process_mode)
         state.append_message(state.roles[0], image_message)
         state.append_message(state.roles[1], "▌")
         state.skip_next = False
@@ -130,9 +162,8 @@ def send_data(state, image, image_process_mode, text_caption, task, request: gr.
             state.skip_next = True
             return (state, None, None, image) + (no_change_btn,) * 7
 
-        # Reset the conversation state for new text inputs too
+        # Reset conversation for text inputs too
         state = default_conversation.copy()
-        
         state.append_message(state.roles[0], text_caption)
         state.append_message(state.roles[1], "▌")
         state.skip_next = False