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