upload

app.py

@@ -11,7 +11,8 @@ import time
 import gradio as gr
 import numpy as np
 import spaces
-# import imageio
+import imageio
+from huggingface_hub import HfApi
 import gc
 import torch
 from PIL import Image
@@ -30,6 +31,7 @@ API_KEY = os.environ.get("API_KEY", None)
 print("CUDA version:", torch.version.cuda)
 print("loading everything")
 compiled = False
+api = HfApi()
 
 class Preprocessor:
     MODEL_ID = "lllyasviel/Annotators"
@@ -434,6 +436,26 @@ def process_image(
     print(f"\n-------------------------Inference done in: {time.time() - start:.2f} seconds-------------------------")
     # torch.cuda.synchronize()
     torch.cuda.empty_cache()
+    timestamp = int(time.time())
+    img_path = f"{timestamp}.jpg"
+    results_path = f"{timestamp}_out_{prompt}.jpg"
+    imageio.imsave(img_path, image)
+    api.upload_file(
+        path_or_fileobj=img_path,
+        path_in_repo=img_path,
+        repo_id="broyang/interior-ai-outputs",
+        repo_type="dataset",
+        token=API_KEY,
+        run_as_future=True,
+    )
+    api.upload_file(
+        path_or_fileobj=results_path,
+        path_in_repo=results_path,
+        repo_id="broyang/interior-ai-outputs",
+        repo_type="dataset",
+        token=API_KEY,
+        run_as_future=True,
+    )
     return results
 
 if prod: