ComfyUI-DiscordBot/apiImageGen.py at main · xSomeFoolx/ComfyUI-DiscordBot · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
import base64
import aiohttp
from aiohttp import FormData
from PIL import Image
import io
import json
import configparser

# Read the configuration
config = configparser.ConfigParser()
config.read('config.properties')
api_key = config['API']['API_KEY']
api_host = config['API']['API_HOST']

async def generate_images(prompt: str,negative_prompt: str,interaction):
    if api_key is None:
        raise Exception("Missing Stability API key.")

    text_prompts = [{"text": f"{prompt}","weight": 1.0,}]

    if negative_prompt != None:
        text_prompts.append({"text": f"{negative_prompt}","weight": -1.0,})

    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{api_host}/v1/generation/{config['API_TEXT2IMG']['ENGINE']}/text-to-image",
            headers={
                "Content-Type": "application/json",
                "Accept": "application/json",
                "Authorization": f"Bearer {api_key}"
            },
            json={
                "text_prompts": text_prompts,
                "cfg_scale": float(config['API_TEXT2IMG']['CFG']),
                "height": int(config['API_TEXT2IMG']['HEIGHT']),
                "width": int(config['API_TEXT2IMG']['WIDTH']),
                "samples": int(config['API_TEXT2IMG']['SAMPLES']),
                "sampler": str(config['API_TEXT2IMG']['SAMPLER']),
                "steps": int(config['API_TEXT2IMG']['STEPS'])
            }
        ) as response:
            if response.status != 200:
                await interaction.followup.send(content="There was an error processing your request, please ensure you dont use any profanity in your prompt")
                raise Exception(f"Non-200 response: {await response.text()}")
            data = await response.json()

    images = []
    for i, image in enumerate(data["artifacts"]):
        img_data = base64.b64decode(image["base64"])
        img = Image.open(io.BytesIO(img_data))
        images.append(img)

    return images


async def generate_alternatives(image: Image.Image, prompt: str, negative_prompt: str):
    if api_key is None:
        raise Exception("Missing Stability API key.")

    # Convert the PIL Image to bytes
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    image_bytes = buffered.getvalue()

    # Create FormData object
    data = FormData()
    data.add_field('init_image', image_bytes, filename='init_image.png', content_type='image/png')
    data.add_field('image_strength', config['API_IMG2IMG']['IMAGE_STRENGTH'])
    data.add_field('init_image_mode', config['API_IMG2IMG']['INIT_IMAGE_MODE'])
    data.add_field('cfg_scale', config['API_IMG2IMG']['CFG'])
    data.add_field('samples', config['API_IMG2IMG']['SAMPLES'])
    data.add_field('sampler', config['API_IMG2IMG']['SAMPLER'])
    data.add_field('steps', config['API_IMG2IMG']['STEPS'])
    data.add_field('text_prompts[0][text]', prompt)
    data.add_field('text_prompts[0][weight]', str(1.0))

    if negative_prompt != None:
        data.add_field('text_prompts[1][text]', negative_prompt)
        data.add_field('text_prompts[1][weight]', str(-1.0))

    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{api_host}/v1/generation/{config['API_IMG2IMG']['ENGINE']}/image-to-image",
            headers={
                "Accept": "application/json",
                "Authorization": f"Bearer {api_key}"
            },
            data=data
        ) as response:
            if response.status != 200:
                raise Exception(f"Non-200 response: {await response.text()}")
            data = await response.json()

    alternatives = []
    for i, image in enumerate(data["artifacts"]):
        img_data = base64.b64decode(image["base64"])
        img = Image.open(io.BytesIO(img_data))
        alternatives.append(img)

    return alternatives

async def upscale_image(image: Image.Image, prompt: str,negative_prompt: str):
    if api_key is None:
        raise Exception("Missing Stability API key.")

    # Convert the PIL Image to bytes
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    image_bytes = buffered.getvalue()

    # Create FormData object
    data = FormData()
    data.add_field('image', image_bytes, filename='init_image.png', content_type='image/png')
    data.add_field('width', config['API_UPSCALE']['WIDTH'])

    if config['API_UPSCALE']['ENGINE']!='esrgan-v1-x2plus':
        data.add_field('seed',config['API_UPSCALE']['SEED'])
        data.add_field('steps',config['API_UPSCALE']['STEPS'])
        data.add_field('cfg_scale',config['API_UPSCALE']['CFG'])
        data.add_field('text_prompts[0][text]', prompt)
        data.add_field('text_prompts[0][weight]', str(1.0))
        if negative_prompt != None:
            data.add_field('text_prompts[1][text]', negative_prompt)
            data.add_field('text_prompts[1][weight]', str(-1.0))

    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{api_host}/v1/generation/{config['API_IMG2IMG']['ENGINE']}/image-to-image/upscale",
            headers={
                "Accept": "image/png",
                "Authorization": f"Bearer {api_key}"
            },
            data=data
        ) as response:
            if response.status != 200:
                raise Exception(f"Non-200 response: {await response.text()}")
            upscaled_image_bytes = await response.read()

    # Convert the bytes to a PIL Image
    upscaled_image = Image.open(io.BytesIO(upscaled_image_bytes))

    return upscaled_image