matrix-pygmalion-bot/matrix_pygmalion_bot/ai/stablehorde.py

import asyncio
import os, tempfile
import logging

import requests

from transformers import AutoTokenizer, AutoConfig
from huggingface_hub import hf_hub_download

from .pygmalion_helpers import get_full_prompt, num_tokens
#from .llama_helpers import get_full_prompt, num_tokens

logger = logging.getLogger(__name__)


async def generate_sync(
    prompt: str,
    api_key: str,
    bot_name: str,
):
    # Set the API endpoint URL
    endpoint = "https://koboldai.net/api/v2/generate/async"

    # Set the headers for the request
    headers = {
        "Content-Type": "application/json",
        "accept": "application/json",
        "apikey": f"{api_key}"
    }

    max_new_tokens = 200
    prompt_num_tokens = await num_tokens(prompt)

    # Define your inputs
    input_data = {
        "prompt": prompt,
        "params": {
            "n": 1,
#            "frmtadsnsp": False,
#            "frmtrmblln": False,
#            "frmtrmspch": False,
#            "frmttriminc": False,
            "max_context_length": 1024,
            "max_length": 512,
            "rep_pen": 1.1,
            "rep_pen_range": 1024,
            "rep_pen_slope": 0.7,
#            "singleline": False,
#            "soft_prompt": "",
            "temperature": 0.75,
            "tfs": 1.0,
            "top_a": 0.0,
            "top_k": 0,
            "top_p": 0.9,
            "typical": 1.0,
#            "sampler_order": [0],
        },
        "softprompts": [],
        "trusted_workers": False,
        "nsfw": True,
#        "workers": [],
        "models": ["PygmalionAI/pygmalion-6b"]
    }


    logger.info(f"sending request to koboldai.net")

    # Make the request
    r = requests.post(endpoint, json=input_data, headers=headers, timeout=180)

    r_json = r.json()
    logger.info(r_json)
    status = r_json["message"]

    if "id" in r_json:
        job_id = r_json["id"]
        TIMEOUT = 360
        DELAY = 11
        for i in range(TIMEOUT//DELAY):
            endpoint = "https://koboldai.net/api/v2/generate/text/status/" + job_id
            r = requests.get(endpoint, headers=headers)
            r_json = r.json()
            logger.info(r_json)
            if "done" not in r_json:
                return "<ERROR>"
            if r_json["done"] == True:
                text = r_json["generations"][0]["text"]
                answer = text.removeprefix(prompt)
                idx = answer.find(f"\nYou:")
                if idx != -1:
                    reply = answer[:idx].strip()
                else:
                    reply = answer.removesuffix('<|endoftext|>').strip()
                reply = reply.replace(f"\n{bot_name}: ", " ")
                reply.replace(f"\n<BOT>: ", " ")
                return reply
            else:
                await asyncio.sleep(DELAY)
    else:
        return f"<ERROR> {status}"


async def generate_image(input_prompt: str, negative_prompt: str, model: str, api_key: str):

    # Set the API endpoint URL
    endpoint = "https://stablehorde.net/api/v2/generate/async"

    # Set the headers for the request
    headers = {
        "Content-Type": "application/json",
        "accept": "application/json",
        "apikey": f"{api_key}"
    }

    # Define your inputs
    input_data = {
        "prompt": input_prompt,
        "params": {
#            "negative_prompt": negative_prompt,
            "width": 512,
            "height": 512,
        },
        "nsfw": True,
        "trusted_workers": False,
#        "workers": [],
        "models": [f"{model}"]
    }

    logger.info(f"sending request to stablehorde.net")

    # Make the request
    r = requests.post(endpoint, json=input_data, headers=headers)
    r_json = r.json()
    logger.info(r_json)

    if r.status_code == 202:
        #status = r_json["message"]
        job_id = r_json["id"]
        TIMEOUT = 360
        DELAY = 11
        output = None
        for i in range(TIMEOUT//DELAY):
            endpoint = "https://stablehorde.net/api/v2/generate/status/" + job_id
            r = requests.get(endpoint, headers=headers)
            r_json = r.json()
            logger.info(r_json)
            #status = r_json["message"]
            if "done" not in r_json:
                return "<ERROR>"
            if "faulted" in r_json and r_json["faulted"] == True:
                return "<ERROR>"
            if r_json["done"] == True:
                output = r_json["generations"]
                break
            else:
                await asyncio.sleep(DELAY)

        if not output:
            raise ValueError(f"<ERROR>")

        os.makedirs("./images", exist_ok=True)
        files = []
        for image in output:
            temp_name = next(tempfile._get_candidate_names())
            filename = "./images/" + temp_name + ".jpg"
            await download_image(image["img"], filename)
            files.append(filename)

        return files

async def generate_image1(input_prompt: str, negative_prompt: str, api_key: str):
    return await generate_image(input_prompt, negative_prompt, "Deliberate", api_key)

async def generate_image2(input_prompt: str, negative_prompt: str, api_key: str):
    return await generate_image(input_prompt, negative_prompt, "PFG", api_key)

async def generate_image3(input_prompt: str, negative_prompt: str, api_key: str):
    return await generate_image(input_prompt, negative_prompt, "Hassanblend", api_key)

async def download_image(url, path):
    r = requests.get(url, stream=True)
    if r.status_code == 200:
        with open(path, 'wb') as f:
            for chunk in r:
                f.write(chunk)