Create engine.py

engine.py

@@ -0,0 +1,199 @@
+import os
+import time
+import requests
+import random
+import json
+import base64
+from io import BytesIO
+from PIL import Image
+
+
+class Prodia:
+    def __init__(self, api_key, base=None):
+        self.base = base or "https://api.prodia.com/v1"
+        self.headers = {
+            "X-Prodia-Key": api_key
+        }
+
+    def sd_controlnet(self, params):
+        response = self._post(f"{self.base}/sd/controlnet", params)
+        return response.json()
+
+    def sd_transform(self, params):
+        response = self._post(f"{self.base}/sd/transform", params)
+        return response.json()
+
+    def sd_generate(self, params):
+        response = self._post(f"{self.base}/sd/generate", params)
+        return response.json()
+
+    def sdxl_generate(self, params):
+        response = self._post(f"{self.base}/sdxl/generate", params)
+        return response.json()
+
+    def upscale_image(self, params):
+        response = self._post(f"{self.base}/upscale", 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'])
+
+        if job_result['status'] == 'failed':
+            raise Exception("Job failed")
+
+        return job_result
+
+    def upload(self, file):
+        files = {'file': open(file, 'rb')}
+        img_id = requests.post(os.getenv("IMAGES_1"), files=files).json()['id']
+
+        payload = {
+            "content": "",
+            "nonce": f"{random.randint(1, 10000000)}H9X42KSEJFNNH",
+            "replies": [],
+            "attachments":
+                [img_id]
+        }
+        resp = requests.post(os.getenv("IMAGES_2"), json=payload, headers={"x-session-token": os.getenv("session-token")})
+        return f"{os.getenv('IMAGES_1')}/{img_id}/{resp.json()['attachments'][0]['filename']}"
+
+    def list_models(self):
+        response = self._get(f"{self.base}/models/list")
+        return response.json()
+
+    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
+
+
+def image_to_base64(image_path):
+    # Open the image with PIL
+    with Image.open(image_path) as image:
+        # Convert the image to bytes
+        buffered = BytesIO()
+        image.save(buffered, format="PNG")  # You can change format to PNG if needed
+
+        # Encode the bytes to base64
+        img_str = base64.b64encode(buffered.getvalue())
+
+    return img_str.decode('utf-8')  # Convert bytes to string
+
+
+prodia_client = Prodia(api_key=os.getenv("PRODIA_X_KEY"))
+
+
+def generate_sdxl(prompt, negative_prompt, model, steps, sampler, cfg_scale, seed):
+    result = prodia_client.sdxl_generate({
+        "prompt": prompt,
+        "negative_prompt": negative_prompt,
+        "model": model,
+        "steps": steps,
+        "sampler": sampler,
+        "cfg_scale": cfg_scale,
+        "seed": seed
+    })
+
+    job = prodia_client.wait(result)
+
+    return job["imageUrl"]
+
+
+def generate_sd(prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed, upscale):
+    result = prodia_client.sd_generate({
+        "prompt": prompt,
+        "negative_prompt": negative_prompt,
+        "model": model,
+        "steps": steps,
+        "sampler": sampler,
+        "cfg_scale": cfg_scale,
+        "seed": seed,
+        "upscale": upscale,
+        "width": width,
+        "height": height
+    })
+
+    job = prodia_client.wait(result)
+
+    return job["imageUrl"]
+
+
+def transform_sd(image, model, prompt, denoising_strength, negative_prompt, steps, cfg_scale, seed, upscale, sampler):
+    image_url = prodia_client.upload(image)
+    result = prodia_client.sd_transform({
+        "imageUrl": image_url,
+        'model': model,
+        'prompt': prompt,
+        'denoising_strength': denoising_strength,
+        'negative_prompt': negative_prompt,
+        'steps': steps,
+        'cfg_scale': cfg_scale,
+        'seed': seed,
+        'upscale': upscale,
+        'sampler': sampler
+    })
+
+    job = prodia_client.wait(result)
+
+    return job["imageUrl"]
+
+
+def controlnet_sd(image, controlnet_model, controlnet_module, threshold_a, threshold_b, resize_mode, prompt, negative_prompt, steps, cfg_scale, seed, sampler, width, height):
+    image_url = prodia_client.upload(image)
+    result = prodia_client.sd_transform({
+        "imageUrl": image_url,
+        "controlnet_model": controlnet_model,
+        "controlnet_module": controlnet_module,
+        "threshold_a": threshold_a,
+        "threshold_b": threshold_b,
+        "resize_mode": int(resize_mode),
+        "prompt": prompt,
+        'negative_prompt': negative_prompt,
+        'steps': steps,
+        'cfg_scale': cfg_scale,
+        'seed': seed,
+        'sampler': sampler,
+        "height": height,
+        "width": width
+    })
+
+    job = prodia_client.wait(result)
+
+    return job["imageUrl"]
+
+def image_upscale(image, scale_by):
+    image_url = prodia_client.upload(image)
+    result = prodia_client.upscale_image({
+        'imageUrl': image_url,
+        'resize': int(scale_by)
+    })
+
+    job = prodia_client.wait(result)
+
+    return job["imageUrl"]
+
+def get_models():
+    return prodia_client.list_models()