Update app.py

app.py

@@ -1,5 +1,6 @@
 # app.py
 import gradio as gr
+from gradio import Request
 import requests
 import os
 import time
@@ -13,10 +14,10 @@ ASTRIA_API_KEY = os.environ.get("ASTRIA_API_KEY")
 
 # --- Core Application Logic ---
 
-def virtual_tryon(human_img_path, garment_img_path, garment_type):
+def virtual_tryon(human_img_path, garment_img_path, garment_type, request: Request):
     """
     This function handles the entire process:
-    1. Validates inputs and API key.
+    1. Validates inputs, API key, and user quota by using the request object.
     2. Constructs and sends a request to create a Tune and associated Prompt.
     3. Polls the Prompt endpoint until the model is trained and image is generated.
     4. Yields status updates to the Gradio UI.
@@ -29,20 +30,56 @@ def virtual_tryon(human_img_path, garment_img_path, garment_type):
     if not human_img_path or not garment_img_path:
         raise gr.Error("Please upload both a human image and a garment image.")
 
-    yield None, "Step 1/3: Preparing and uploading your images..."
+    user_ip = request.client.host if request else "unknown"
+    if user_ip == "unknown":
+        # This error is not shown to the user but helps in debugging
+        print("Warning: Could not determine user IP.")
+        # Proceed without a quota check if IP is not found
+        pass
+    else:
+        headers = {"Authorization": f"Bearer {ASTRIA_API_KEY}"}
+        tune_title = f"Gradio VTO {user_ip}"
+
+        # 2. Usage Quota Check
+        yield None, f"Step 1/4: Checking usage quota..."
+        try:
+            # Query the API for tunes created by this user IP
+            quota_check_url = f"{ASTRIA_API_BASE_URL}/tunes"
+            params = {"title": tune_title}
+            quota_response = requests.get(quota_check_url, headers=headers, params=params)
+            quota_response.raise_for_status()
+            user_tunes = quota_response.json()
+
+            # Enforce a limit of 5 trials
+            if len(user_tunes) > 4:
+                raise gr.Error(
+                    f"Usage quota reached (5 trials). To continue, please use the API directly at astria.ai."
+                )
+
+            trials_used = len(user_tunes)
+            yield None, f"Quota check passed ({trials_used}/5 trials used). Proceeding to step 2..."
+
+        except requests.exceptions.RequestException as e:
+            raise gr.Error(f"API request failed during quota check: {e}")
+        except Exception as e:
+            # Catch the gr.Error from the check and re-raise it
+            if isinstance(e, gr.Error):
+                raise e
+            raise gr.Error(f"An unexpected error occurred during quota check: {e}")
+
 
-    headers = {"Authorization": f"Bearer {ASTRIA_API_KEY}"}
+    yield None, "Step 2/4: Preparing and uploading your images..."
     prompt_id = None
 
     try:
-        # 2. Construct and Send API Request to create Tune and Prompt
+        # 3. Construct and Send API Request to create Tune and Prompt
         with open(human_img_path, "rb") as human_f, open(garment_img_path, "rb") as garment_f:
             files = [
                 ("tune[images][]", ("garment.jpg", garment_f.read(), "image/jpeg")),
                 ("tune[prompts_attributes][][input_image]", ("human.jpg", human_f.read(), "image/jpeg")),
             ]
             data = {
-                "tune[title]": f"Gradio VTO - {int(time.time())}",
+                "tune[title]": f"Gradio VTO - {user_ip if user_ip != 'unknown' else 'anonymous'}",
                 "tune[name]": garment_type,
                 "tune[model_type]": "faceid",
                 "tune[base_tune_id]": "1504944", # This is the base model for Flux although no image generation is done in this case
@@ -72,9 +109,9 @@ def virtual_tryon(human_img_path, garment_img_path, garment_type):
     except Exception as e:
         raise gr.Error(f"An unexpected error occurred during upload: {e}")
 
-    yield None, f"Step 2/3: Task created (Prompt ID: {prompt_id}). Waiting for model training and generation..."
+    yield None, f"Step 3/4: Task created (Prompt ID: {prompt_id}). Waiting for model training and generation..."
 
-    # 3. Polling the Prompt Endpoint for Results
+    # 4. Polling the Prompt Endpoint for Results
     start_time = time.time()
     timeout = 360  # 6-minute timeout for training and generation
 
@@ -90,13 +127,13 @@ def virtual_tryon(human_img_path, garment_img_path, garment_type):
                 images = poll_data.get("images", [])
                 if images and images[0]:
                     output_image_url = images[0]
-                    yield output_image_url, "Step 3/3: Generation successful!"
+                    yield output_image_url, "Step 4/4: Generation successful!"
                     return  # End the function
                 else:
                     # This case handles when training is done but image isn't ready yet.
-                    yield None, "Step 2/3: Model trained. Finalizing image..."
+                    yield None, "Step 3/4: Model trained. Finalizing image..."
             else:
-                yield None, "Step 2/3: Model is training. Checking again in 10 seconds..."
+                yield None, "Step 3/4: Model is training. Checking again in 10 seconds..."
 
             time.sleep(10)  # Wait before polling again
 
@@ -127,6 +164,9 @@ ASTRIA_API_KEY = os.environ.get("ASTRIA_API_KEY") # Recommended: Load from envir
 human_img_path = "path/to/your/human.jpg"
 garment_img_path = "path/to/your/garment.jpg"
 garment_type = "shirt"  # Can be "shirt", "pants", or "dress"
+# To track usage, you can include a user identifier in the title
+user_identifier = "user-ip-or-id" 
+tune_title = f"My App VTO - {user_identifier} - {int(time.time())}"
 
 headers = {"Authorization": f"Bearer {ASTRIA_API_KEY}"}
 
@@ -139,7 +179,7 @@ with open(human_img_path, "rb") as human_f, open(garment_img_path, "rb") as garm
         ("tune[prompts_attributes][][input_image]", ("human.jpg", human_f.read(), "image/jpeg")),
     ]
     data = {
-        "tune[title]": f"My VTO API Test - {int(time.time())}",
+        "tune[title]": tune_title,
         "tune[name]": garment_type,
         "tune[model_type]": "faceid",
         "tune[base_tune_id]": "1504944", # VTO Base Model ID
@@ -157,66 +197,4 @@ print(response.json())
 # See this Space's app.py for a full polling example.
 """
 
-with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
-    # --- Banner and Main Title ---
-    gr.HTML(
-        """
-        <div style="text-align: center; padding: 20px; background-color: #f0f8ff; border-radius: 10px; margin-bottom: 20px;">
-            <h2 style="color: #333;">Powered by the Astria.ai API</h2>
-            <p style="color: #555; font-size: 1.1em;">This Virtual Try-On demo uses a powerful, single-call API endpoint.
-            <br>
-            You can integrate this functionality directly into your own applications.
-            </p>
-            <a href="https://docs.astria.ai/docs/use-cases/virtual-try-on/" target="_blank" style="text-decoration: none;">
-                <button style="padding: 10px 20px; font-size: 1em; color: white; background-color: #007bff; border: none; border-radius: 5px; cursor: pointer;">
-                    Read the Full API Documentation
-                </button>
-            </a>
-        </div>
-        """
-    )
-    gr.Markdown("# Virtual Fashion Try-On")
-    gr.Markdown(
-        """
-        **Instructions:** Upload a photo of a person and a photo of a piece of clothing to see them wear it.
-        
-        This space is made available through the Astria API, which allows you to create virtual try-on experiences with just a few lines of code. 
-        Astria fine-tuning API allows you to create custom AI models from your own images, and then use Virtual Try-On to make those models wear different clothing items.
-        This can be used for fashion e-commerce, virtual fitting rooms, music festivals and concerts and sport events activation campaigns. See more in [Astria Virtual Try-On documentation](https://docs.astria.ai/docs/use-cases/virtual-try-on/).
-        """
-    )
-
-    # --- Main Application UI ---
-    with gr.Row():
-        with gr.Column(scale=1):
-            human_image = gr.Image(type="filepath", label="Human Image", height=300)
-            garment_image = gr.Image(type="filepath", label="Garment Image", height=300)
-            garment_type = gr.Dropdown(
-                ["shirt", "pants", "dress"],
-                label="Garment Type",
-                info="Select the type of clothing item you uploaded.",
-                value="shirt"
-            )
-            submit_btn = gr.Button("Generate", variant="primary")
-        with gr.Column(scale=1):
-            status_text = gr.Textbox(label="Status", interactive=False, lines=2, placeholder="Upload images and click Generate...")
-            result_image = gr.Image(label="Result", height=615, interactive=False)
-
-    # --- API Usage Code Snippet ---
-    with gr.Accordion("Show API Usage Code (Python)", open=False):
-        gr.Code(
-            value=code_snippet,
-            language="python",
-            label="Python Request Example",
-            interactive=False
-        )
-
-    # --- Event Listener ---
-    submit_btn.click(
-        fn=virtual_tryon,
-        inputs=[human_image, garment_image, garment_type],
-        outputs=[result_image, status_text]
-    )
-
-if __name__ == "__main__":
-    demo.launch()
+with gr.Blocks(