Tiktok-Talent-Info/locustfile.py

import os
import time
import asyncio
from locust import FastHttpUser, HttpUser, task, between, constant

import torch
print(torch.cuda.is_available())  

class FastAPIUser(FastHttpUser):
    # wait_time = between(1, 3)  # Wait time between requests
    wait_time = constant(0)  # send request simultaneously
    
    def ensure_event_loop(self):
        """Ensures each Locust task runs in a fresh asyncio event loop."""
        try:
            asyncio.get_running_loop()
        except RuntimeError:
            asyncio.set_event_loop(asyncio.new_event_loop())

    # @task
    # def text_query(self):
    #     """Simulate text query request."""
    #     url = "/api/text"
        
    #     # start_event.wait()  # Synchronize all users

    #     start_time = time.time()
    #     response = self.client.post(url, data={"question": "What is FastAPI?"})
    #     total_time = time.time() - start_time

    #     if response.status_code == 200:
    #         print(f"Success: {response.json()} (Time: {total_time:.2f}s)")
    #     else:
    #         print(f"Error: {response.status_code} - {response.text}")

    # @task
    # def upload_image(self):
    #     """Sends an image with the correct query format."""
    #     url = "/api/image"
    #     image_path = "../profile/1.png"

    #     if not os.path.exists(image_path):
    #         print(f"Error: File {image_path} not found!")
    #         return

    #     query_text = (
    #         "Extract the following information from this image and return the result in JSON format:\n"
    #         "- Name: <name>\n"
    #         "- ID: <id>\n"
    #         "- Profile Picture: <url>\n"
    #         "- Follower Count: <count>\n"
    #         "- Likes Count: <count>\n"
    #         "- Bio: <bio>\n"
    #         "- Following Count: <count>\n"
    #         "- External Links: <links>\n"
    #         "Do not include any disclaimers or comments like 'I'm sorry' or 'I can't assist'."
    #     )

    #     with open(image_path, "rb") as image_file:
    #         files = {
    #             "file": ("1.png", image_file, "image/png"),
    #             "question": (None, query_text)
    #         }

    #         start_time = time.time()
    #         response = self.client.post(url, files=files)
    #         total_time = time.time() - start_time

    #         if response.status_code == 200:
    #             print(f"Success: {response.json()} (Time: {total_time:.2f}s)")
    #         else:
    #             print(f"Error: {response.status_code} - {response.text}")

    @task
    def video_query(self):
        """Uploads a video and asks a detailed question."""
        url = "/api/video"
        video_path = "../video/1.1.mp4"

        if not os.path.exists(video_path):
            print(f"Error: File {video_path} not found!")
            return

        with open(video_path, "rb") as file:
            files = {
                "file": ("1.1.mp4", file, "video/mp4"),
                "question": (None, """Based on the given images and audio script, extract detailed 
                information about the products recommended in the video and format the output as JSON with 
                the following fields:
                1. **Product Name**: The specific name of the product, if mentioned.
                2. **Category**: The specific category of the product (e.g., electronics, skincare, casual wear, etc.).
                3. **Styles or Variants**: Any styles, designs, or variants of the product described 
                   (e.g., colors, patterns, sizes, or other distinguishing attributes).
                4. **Highlights**: The unique selling points or notable features emphasized by the anchor 
                   (e.g., benefits, quality, or standout features).
                5. **Promotional Details**: Any additional promotional information mentioned, such as discounts, 
                   offers, or key features that set the product apart."""
                )
            }

            print("Starting video query...")
            start_time = time.time()
            response = self.client.post(url, files=files) 
            total_time = time.time() - start_time

            if response.status_code == 200:
                print(f"Success: {response.json()} (Time: {total_time:.2f}s)")
            else:
                print(f"Error: {response.status_code} - {response.text}")

            print("Video query completed.")

# locust -f locustfile.py

# curl -X POST http://localhost:8002/api/image \
#   -F "file=@../profile/1.png" \
#   -F "question=What is this image about?"