engine.py

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()