Upload 2 files

const.py

@@ -0,0 +1,17 @@
+SAMPLER_LIST = ["Euler", "Euler a", "LMS", "Heun", "DPM2", "DPM2 a", "DPM++ 2S a", "DPM++ 2M", "DPM++ SDE", "DPM fast",
+                "DPM adaptive", "LMS Karras", "DPM2 Karras", "DPM2 a Karras", "DPM++ 2S a Karras", "DPM++ 2M Karras",
+                "DPM++ SDE Karras", "DDIM", "PLMS"]
+
+SDXL_MODEL_LIST = ["sd_xl_base_1.0.safetensors [be9edd61]", "dreamshaperXL10_alpha2.safetensors [c8afe2ef]",
+                   "dynavisionXL_0411.safetensors [c39cc051]"]
+
+CMODULES = ['none', 'canny', 'depth', 'depth_leres', 'depth_leres++', 'hed', 'hed_safe', 'mediapipe_face', 'mlsd',
+           'normal_map', 'openpose', 'openpose_hand', 'openpose_face', 'openpose_faceonly', 'openpose_full',
+           'clip_vision', 'color', 'pidinet', 'pidinet_safe', 'pidinet_sketch', 'pidinet_scribble', 'scribble_xdog',
+           'scribble_hed', 'segmentation', 'threshold', 'depth_zoe', 'normal_bae', 'oneformer_coco', 'oneformer_ade20k',
+           'lineart', 'lineart_coarse', 'lineart_anime', 'lineart_standard', 'shuffle', 'tile_resample', 'invert',
+           'lineart_anime_denoise', 'reference_only', 'reference_adain', 'reference_adain+attn', 'inpaint',
+           'inpaint_only', 'inpaint_only+lama', 'tile_colorfix', 'tile_colorfix+sharp']
+
+CMODELS = ['control_v11p_sd15_canny [d14c016b]', 'control_v11p_sd15_openpose [cab727d4]',
+           'control_v11p_sd15_softedge [a8575a2a]', 'control_v11p_sd15_scribble [d4ba51ff]']

engine.py

@@ -0,0 +1,188 @@
+import os
+import time
+import requests
+import random
+import json
+import base64
+from io import BytesIO
+from PIL import Image
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+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 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):
+    print(image)
+    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 get_models():
+    return prodia_client.list_models()