Refactor request code to reduce duplication

Image/nodes.py

@@ -1,4 +1,4 @@
-from ..session import session
+from ..session import post
 from io import BytesIO
 from PIL import Image
 import numpy as np
@@ -21,9 +21,7 @@ class TextToImage:
 
     def inference(self, endpoint, prompt):
         payload = {"inputs": prompt}
-        response = session.post(endpoint, json=payload)
-        if response.status_code != 200:
-            raise Exception(response.text)
+        response = post(endpoint, json=payload)
         
         result = BytesIO(response.content)
         image = Image.open(result)
@@ -48,9 +46,7 @@ class Classification:
     TITLE = "HF Image Classification"
 
     def inference(self, endpoint, image):
-        response = session.post(endpoint, data=image)
-        if response.status_code != 200:
-            raise Exception(response.text)
+        response = post(endpoint, data=image)
         result = response.json()
         return {"ui": {"text": result}}
     
@@ -70,9 +66,7 @@ class ObjectDetection:
     TITLE = "HF Image Object Detection"
 
     def inference(self, endpoint, image):
-        response = session.post(endpoint, data=image)
-        if response.status_code != 200:
-            raise Exception(response.text)
+        response = post(endpoint, data=image)
         result = response.json()
         return {"ui": {"text": result}}
     
@@ -92,9 +86,7 @@ class Segmentation:
     TITLE = "HF Image Segmentation"
 
     def inference(self, endpoint, image):
-        response = session.post(endpoint, data=image)
-        if response.status_code != 200:
-            raise Exception(response.text)
+        response = post(endpoint, data=image)
         result = response.json()
         return {"ui": {"text": result}}
 

Text/nodes.py

@@ -1,5 +1,5 @@
 import torch
-from ..session import session
+from ..session import post
 
 class Generation:
     @classmethod
@@ -20,9 +20,7 @@ class Generation:
         json = {
             'inputs': text,
         }
-        response = session.post(endpoint, json=json)
-        if response.status_code != 200:
-            raise Exception(response.text)
+        response = post(endpoint, json=json)
         result = response.json()
         generated = ''.join(x['generated_text'] for x in result)
         return generated 
@@ -46,9 +44,7 @@ class Translation:
         json = {
             'inputs': text,
         }
-        response = session.post(endpoint, json=json)
-        if response.status_code != 200:
-            raise Exception(response.text)
+        response = post(endpoint, json=json)
         result = response.json()
         translation = ''.join(x['translation_text'] for x in result)
         return translation 
@@ -76,9 +72,7 @@ class QuestionAnswering:
                 'context': context,
             },
         }
-        response = session.post(endpoint, json=json)
-        if response.status_code != 200:
-            raise Exception(response.text)
+        response = post(endpoint, json=json)
         result = response.json()
         answer = result['answer']
         return answer
@@ -102,9 +96,7 @@ class FeatureExtraction:
         json = {
             'inputs': text,
         }
-        response = session.post(endpoint, json=json)
-        if response.status_code != 200:
-            raise Exception(response.text)
+        response = post(endpoint, json=json)
         result = response.json()
         cond = torch.tensor(result, dtype=torch.float16).to('cuda')
         return ([[cond, {}]],)

session.py

@@ -1,10 +1,28 @@
 """Session handler"""
 import os
 import requests
+from time import sleep
 session = requests.Session()
 if "HF_AUTH_TOKEN" in os.environ:
     session.headers.update({
         "Authorization": f"Bearer {os.environ['HF_AUTH_TOKEN']}",
     })
 else:
-    print("No 'HF_AUTH_TOKEN' set.")
+    print("No 'HF_AUTH_TOKEN' set.")
+
+def post(url, **kwargs):
+
+    if not url.startswith('http'):
+        url = f'https://api-inference.huggingface.co/models/{url}'
+
+    response = session.post(url, **kwargs)
+    
+    if response.status_code != 200:
+        if 'estimated_time' in response.text:
+            estimated_time = response.json()['estimated_time']
+            model_path = '/'.join(url.split('/')[-2:])
+            print('Waiting for ', estimated_time, ' to load ', model_path)
+            sleep(estimated_time)
+            return post(url, **kwargs)
+        raise Exception(response.text)
+    return response