Update app.py

app.py

@@ -7,6 +7,52 @@ import random
 import os
 import base64
 
+
+class Prodia:
+    def __init__(self, api_key=os.getenv("PRODIA_API_KEY"), base=None):
+        self.base = base or "https://api.prodia.com/v1"
+        self.headers = {
+            "X-Prodia-Key": api_key
+        }
+    
+    def faceswap(self, params):
+        response = self._post(f"{self.base}/faceswap", params)
+        return response.json()
+    
+    def get_job(self, job_id):
+        response = self._get(f"{self.base}/job/{job_id}")
+        return response.json()
+
+    def wait(self, job):
+        job_result = job
+
+        while job_result['status'] not in ['succeeded', 'failed']:
+            time.sleep(0.25)
+            job_result = self.get_job(job['job'])
+
+        return job_result
+
+    def _post(self, url, params):
+        headers = {
+            **self.headers,
+            "Content-Type": "application/json"
+        }
+        response = requests.post(url, headers=headers, data=json.dumps(params))
+
+        if response.status_code != 200:
+            raise Exception(f"Bad Prodia Response: {response.status_code}")
+
+        return response
+
+    def _get(self, url):
+        response = requests.get(url, headers=self.headers)
+
+        if response.status_code != 200:
+            raise Exception(f"Bad Prodia Response: {response.status_code}")
+
+        return response
+
+
 client = Prodia()
 
 
@@ -17,12 +63,13 @@ def infer(source, target):
     source_url = upload_image(source)
     target_url = upload_image(target)
 
-    job = client.faceswap(source_url=source_url, target_url=target_url)
-    res = client.wait(job, raise_on_fail=False)
-    if res.failed:
-        return
+    job = client.faceswap({
+        "sourceUrl": source_url,
+        "targetUrl": target_url
+    })
+    res = client.wait(job)
 
-    return res.image_url
+    return res['imageUrl']
 
 
 def upload_image(file):