Img Gen Rev3

app.py

@@ -28,6 +28,22 @@ def get_completion(inputs, parameters=None, ENDPOINT_URL=""):
                                 data=json.dumps(data))
     return json.loads(response.content.decode("utf-8"))
 
+def get_img_completion(inputs, parameters=None, ENDPOINT_URL=""):
+    headers = {
+      "Authorization": f"Bearer {hf_api_key}",
+      "Content-Type": "application/json"
+    }   
+    data = { "inputs": inputs }
+    if parameters is not None:
+        data.update({"parameters": parameters})
+    response = requests.request("POST",
+                                ENDPOINT_URL,
+                                headers=headers,
+                                data=json.dumps(data))
+    return response.content
+    # return json.loads(response.content.decode("utf-8"))
+    
+
 #text-to-image
 TTI_ENDPOINT = "https://api-inference.huggingface.co/models/ZB-Tech/Text-to-Image"
 # "http://jupyter-api-proxy.internal.dlai/rev-proxy/huggingface/itt"
@@ -58,8 +74,8 @@ def captioner(image):
     return result[0]['generated_text']
 
 def generate(prompt):
-    output = get_completion(prompt, None, TTI_ENDPOINT)
-    image = Image.open(io.BytesIO(output.content))
+    output = get_img_completion(prompt, None, TTI_ENDPOINT)
+    image = Image.open(io.BytesIO(output))
     return image
     # result_image = base64_to_pil(output.content)
     # return result_image