Update app.py

app.py

@@ -2,30 +2,83 @@ import os
 import numpy as np
 from prodiapy import Prodia
 import gradio as gr
+import json
+import requests
 import base64
+import random
 
-prodia = Prodia()
 STYLE_PRESETS = [None, "3d-model", "analog-film", "anime", "cinematic", "comic-book", "digital-art", "enhance", "fantasy-art", "isometric", "line-art", "low-poly",
                  "neon-punk", "origami", "photographic", "pixel-art", "texture", "craft-clay"]
 MAX_SEED = np.iinfo(np.int32).max
 
 
+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 photomaker(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()
+
+
 def generate_image(upload_images, prompt, negative_prompt, style_preset, steps, cfg_scale, strength, seed, progress=gr.Progress(track_tqdm=True)):
     error_if_no_img(prompt)
 
-    job = prodia.create("/photomaker",
-                        imageData=[file_to_base64(img) for img in upload_images],
-                        prompt=prompt,
-                        negative_prompt=negative_prompt,
-                        style_preset=style_preset,
-                        steps=steps,
-                        cfg_scale=cfg_scale,
-                        strength=strength,
-                        seed=seed if seed != 0 else None
-                        )
-    result = prodia.wait(job)
-
-    return result.image_url
+    job = client.photomaker({
+        "imageData": [file_to_base64(img) for img in upload_images],
+        "prompt": prompt,
+        "negative_prompt": negative_prompt,
+        "style_preset": style_preset,
+        "steps": steps, 
+        "cfg_scale": cfg_scale,
+        "strength": strength,
+        "seed": seed if seed != 0 else random.randint(1, MAX_SEED)
+    })              
+    res = prodia.wait(job)
+
+    if res['status'] == "failed":
+        return
+    
+    return res['imageUrl']
 
 
 def error_if_no_img(prompt):