Update app.py

app.py

@@ -3,51 +3,41 @@ import json
 import base64
 import os
 import requests
-from huggingface_hub import HfApi, HfFolder
+from huggingface_hub import upload_file
 
-HF_TOKEN = os.environ.get("HF_CV_ROBOT_TOKEN")
+HF_TOKEN = os.environ.get("HF_TOKEN")
+HF_DATASET_REPO = "OppaAI/Robot_MCP"  # Replace with your dataset repo
 MODEL = "Qwen/Qwen2.5-VL-7B-Instruct"
-REPO_ID = "OppaAI/Robot_MCP"  # Replace with your HF repo
 
-if not HF_TOKEN:
-    print("ERROR: HF_CV_ROBOT_TOKEN environment variable not set.")
+def process_and_describe(payload: dict):
+    if not HF_TOKEN:
+        return {"error": "HF_TOKEN secret not found in Space settings."}
 
-api = HfApi()
-HfFolder.save_token(HF_TOKEN)
-
-def upload_to_hf(filepath: str):
-    filename = os.path.basename(filepath)
-    # Upload to HF repo
-    api.upload_file(
-        path_or_fileobj=filepath,
-        path_in_repo=filename,
-        repo_id=REPO_ID,
-        repo_type="dataset",
-        token=HF_TOKEN,
-        overwrite=True
-    )
-    # Construct raw URL
-    url = f"https://huggingface.co/datasets/{REPO_ID}/resolve/main/{filename}"
-    return url
-
-def process(payload: dict):
     try:
-        if not HF_TOKEN:
-            return {"error": "Hugging Face token missing."}
-
         robot_id = payload.get("robot_id", "unknown")
         image_b64 = payload["image_b64"]
+        image_bytes = base64.b64decode(image_b64)
 
-        # 1️⃣ Save the image locally
-        tmp_path = "/tmp/tmp.jpg"
-        with open(tmp_path, "wb") as f:
-            f.write(base64.b64decode(image_b64))
-        file_size = os.path.getsize(tmp_path)
+        # 1️⃣ Save temporarily
+        local_tmp_path = "/tmp/uploaded_image.jpg"
+        with open(local_tmp_path, "wb") as f:
+            f.write(image_bytes)
+
+        # 2️⃣ Upload to HF dataset repo
+        path_in_repo = f"images/uploaded_image_{len(image_bytes)}.jpg"
+        upload_file(
+            path_or_fileobj=local_tmp_path,
+            path_in_repo=path_in_repo,
+            repo_id=HF_DATASET_REPO,
+            token=HF_TOKEN,
+            repo_type="dataset"
+        )
+        os.remove(local_tmp_path)
 
-        # 2️⃣ Upload to HF to get URL
-        image_url = upload_to_hf(tmp_path)
+        # 3️⃣ Construct public URL
+        image_url = f"https://huggingface.co/datasets/{HF_DATASET_REPO}/resolve/main/{path_in_repo}"
 
-        # 3️⃣ Call VLM with image URL
+        # 4️⃣ Call VLM
         data = {
             "model": MODEL,
             "messages": [
@@ -76,22 +66,22 @@ def process(payload: dict):
             except Exception as e:
                 vlm_text = f"Failed to parse VLM response: {e}, Response={resp.text}"
 
-        # 4️⃣ Return combined info
         return {
-            "saved": True,
-            "file_path": tmp_path,
-            "file_size_bytes": file_size,
+            "saved_to_hf_hub": True,
+            "repo_id": HF_DATASET_REPO,
+            "path_in_repo": path_in_repo,
+            "image_url": image_url,
+            "file_size_bytes": len(image_bytes),
             "robot_id": robot_id,
-            "vlm_description": vlm_text,
-            "image_url": image_url
+            "vlm_description": vlm_text
         }
 
     except Exception as e:
-        return {"error": str(e)}
+        return {"error": f"Failed to upload/describe image: {str(e)}"}
 
 demo = gr.Interface(
-    fn=process,
-    inputs=gr.JSON(label="Input Payload (Dict format)"),
+    fn=process_and_describe,
+    inputs=gr.JSON(label="Input Payload (Dict format with 'image_b64')"),
     outputs=gr.JSON(label="Reply to Jetson"),
     api_name="predict"
 )