ig_bot/SD.py

61 lines
1.9 KiB
Python
Raw Normal View History

2024-05-25 09:13:33 +02:00
import requests
import os
import uuid
import requests
import os
import uuid
import subprocess
import atexit
class ImageResponse:
def __init__(self, url):
self.url = url
class DataContainer:
def __init__(self, images):
self.data = images
class ImageUrl:
def __init__(self, url):
self.url = url
class ImageData:
def __init__(self, image_urls):
self.data = image_urls
class SD:
def __init__(self, images_dir, api_key=""):
self.endpoint = "http://localhost:8188/get_image_from_workflow"
self.images_dir = images_dir # Correctly define the images_dir attribute here
self.images = self._create_images()
def _create_images(self):
endpoint = self.endpoint
images_dir = self.images_dir
class images:
def generate(self, model, prompt: str, style: str="", negative: str="", size: str="1024x1024", quality="standard", n=1):
width, height = size.split('x')
response = requests.post(endpoint, json={
'positive': prompt,
'negative': "",
'workflow': model,
'width': width,
'height': height,
'n': n
})
if response.status_code == 200:
image_data = response.content
filename = f"{uuid.uuid4()}.png"
file_path = os.path.join(images_dir, filename)
with open(file_path, 'wb') as image_file:
image_file.write(image_data)
url = f"http://localhost:8190/{filename}"
# Return an ImageData object containing a list of ImageUrl objects
return ImageData([ImageUrl(url)])
else:
raise Exception("Failed to generate image: " + response.text)
return images()