Update app.py

app.py

@@ -9,9 +9,10 @@ from qwen_vl_utils import process_vision_info
 from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
 import ast
 import os
-from datetime import datetime
 import numpy as np
 from huggingface_hub import hf_hub_download, list_repo_files
+import boto3
+from botocore.exceptions import NoCredentialsError
 
 # Define constants
 DESCRIPTION = "[ShowUI Demo](https://huggingface.co/showlab/ShowUI-2B)"
@@ -35,8 +36,7 @@ for file in files:
     print(f"Downloaded {file} to {file_path}")
 
 model = Qwen2VLForConditionalGeneration.from_pretrained(
-    "./showui-2b",
-    # "showlab/ShowUI-2B",
+    destination_folder,
     torch_dtype=torch.bfloat16,
     device_map="cpu",
 )
@@ -57,20 +57,43 @@ def draw_point(image_input, point=None, radius=5):
         ImageDraw.Draw(image).ellipse((x - radius, y - radius, x + radius, y + radius), fill='red')
     return image
 
-def array_to_image_path(image_array):
+def array_to_image_path(image_array, session_id):
     """Save the uploaded image and return its path."""
     if image_array is None:
         raise ValueError("No image provided. Please upload an image before submitting.")
     img = Image.fromarray(np.uint8(image_array))
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    filename = f"image_{timestamp}.png"
+    filename = f"image_{session_id}.png"
     img.save(filename)
     return os.path.abspath(filename)
 
+# Function to upload the file to S3
+def upload_to_s3(file_name, bucket, object_name=None):
+    """Upload a file to an S3 bucket."""
+    if object_name is None:
+        object_name = file_name
+
+    # Create an S3 client
+    s3 = boto3.client('s3')
+
+    try:
+        s3.upload_file(file_name, bucket, object_name)
+        print(f"Uploaded {file_name} to {bucket}/{object_name}.")
+        return True
+    except FileNotFoundError:
+        print(f"The file {file_name} was not found.")
+        return False
+    except NoCredentialsError:
+        print("Credentials not available.")
+        return False
+
 @spaces.GPU
 def run_showui(image, query):
     """Main function for inference."""
-    image_path = array_to_image_path(image)
+    session_id = datetime.now().strftime("%Y%m%d_%H%M%S")
+    image_path = array_to_image_path(image, session_id)
+    
+    # Upload the image to S3
+    upload_to_s3(image_path, 'altair.storage', object_name=f"ootb/images/{os.path.basename(image_path)}")
 
     messages = [
         {
@@ -84,9 +107,7 @@ def run_showui(image, query):
     ]
 
     # Prepare inputs for the model
-
     global model
-
     model = model.to("cuda")
     
     text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
@@ -114,11 +135,11 @@ def run_showui(image, query):
 
     # Draw the point on the image
     result_image = draw_point(image_path, click_xy, radius=10)
-    return result_image, str(click_xy)
+    return result_image, str(click_xy), session_id
 
-# Function to record votes
-def record_vote(vote_type, image_path, query, action_generated):
-    """Record a vote in a JSON file."""
+# Modify the record_vote function
+def record_vote(vote_type, image_path, query, action_generated, session_id):
+    """Record a vote in a JSON file and upload to S3."""
     vote_data = {
         "vote_type": vote_type,
         "image_path": image_path,
@@ -126,27 +147,35 @@ def record_vote(vote_type, image_path, query, action_generated):
         "action_generated": action_generated,
         "timestamp": datetime.now().isoformat()
     }
-    with open("votes.json", "a") as f:
+    
+    local_file_name = f"votes_{session_id}.json"
+    
+    # Append vote data to the local JSON file
+    with open(local_file_name, "a") as f:
         f.write(json.dumps(vote_data) + "\n")
+    
+    # Upload the updated JSON file to S3
+    upload_to_s3(local_file_name, 'altair.storage', object_name=f"ootb/votes/{local_file_name}")
+
     return f"Your {vote_type} has been recorded. Thank you!"
 
-# Helper function to handle vote recording
-def handle_vote(vote_type, image_path, query, action_generated):
+# Use session_id in the handle_vote function
+def handle_vote(vote_type, image_path, query, action_generated, session_id):
     """Handle vote recording by using the consistent image path."""
     if image_path is None:
         return "No image uploaded. Please upload an image before voting."
-    return record_vote(vote_type, image_path, query, action_generated)
+    return record_vote(vote_type, image_path, query, action_generated, session_id)
 
 # Load logo and encode to Base64
 with open("./assets/showui.png", "rb") as image_file:
     base64_image = base64.b64encode(image_file.read()).decode("utf-8")
 
-
 # Define layout and UI
 def build_demo(embed_mode, concurrency_count=1):
     with gr.Blocks(title="ShowUI Demo", theme=gr.themes.Default()) as demo:
         # State to store the consistent image path
         state_image_path = gr.State(value=None)
+        state_session_id = gr.State(value=None)
 
         if not embed_mode:
             gr.HTML(
@@ -230,19 +259,35 @@ def build_demo(embed_mode, concurrency_count=1):
                     raise ValueError("No image provided. Please upload an image before submitting.")
                 
                 # Generate consistent image path and store it in the state
-                image_path = array_to_image_path(image)
-                return run_showui(image, query) + (image_path,)
+                result_image, click_coords, session_id = run_showui(image, query)
+                return result_image, click_coords, session_id
+
+            def on_image_upload(image):
+                """Generate a new session ID when a new image is uploaded."""
+                if image is None:
+                    raise ValueError("No image provided. Please upload an image.")
+                
+                # Generate a new session ID
+                new_session_id = datetime.now().strftime("%Y%m%d_%H%M%S")
+                return new_session_id
+
+            imagebox.upload(
+                on_image_upload,
+                inputs=imagebox,
+                outputs=state_session_id,
+                queue=False
+            )
 
             submit_btn.click(
                 on_submit,
                 [imagebox, textbox],
-                [output_img, output_coords, state_image_path],
+                [output_img, output_coords, state_session_id],
             )
 
             clear_btn.click(
                 lambda: (None, None, None, None, None),
                 inputs=None,
-                outputs=[imagebox, textbox, output_img, output_coords, state_image_path],  # Clear all outputs
+                outputs=[imagebox, textbox, output_img, output_coords, state_session_id],  # Clear all outputs
                 queue=False
             )
 
@@ -254,33 +299,34 @@ def build_demo(embed_mode, concurrency_count=1):
 
             # Record vote actions without feedback messages
             vote_btn.click(
-                lambda image_path, query, action_generated: handle_vote(
-                    "upvote", image_path, query, action_generated
+                lambda image_path, query, action_generated, session_id: handle_vote(
+                    "upvote", image_path, query, action_generated, session_id
                 ),
-                inputs=[state_image_path, textbox, output_coords],
+                inputs=[state_image_path, textbox, output_coords, state_session_id],
                 outputs=[],
                 queue=False
             )
 
             downvote_btn.click(
-                lambda image_path, query, action_generated: handle_vote(
-                    "downvote", image_path, query, action_generated
+                lambda image_path, query, action_generated, session_id: handle_vote(
+                    "downvote", image_path, query, action_generated, session_id
                 ),
-                inputs=[state_image_path, textbox, output_coords],
+                inputs=[state_image_path, textbox, output_coords, state_session_id],
                 outputs=[],
                 queue=False
             )
 
             flag_btn.click(
-                lambda image_path, query, action_generated: handle_vote(
-                    "flag", image_path, query, action_generated
+                lambda image_path, query, action_generated, session_id: handle_vote(
+                    "flag", image_path, query, action_generated, session_id
                 ),
-                inputs=[state_image_path, textbox, output_coords],
+                inputs=[state_image_path, textbox, output_coords, state_session_id],
                 outputs=[],
                 queue=False
             )
 
     return demo
+
 # Launch the app
 if __name__ == "__main__":
     demo = build_demo(embed_mode=False)
@@ -289,4 +335,4 @@ if __name__ == "__main__":
         server_port=7860,
         ssr_mode=False,
         debug=True,
-    )
+    )