61 lines
1.9 KiB
Python
61 lines
1.9 KiB
Python
|
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()
|
||
|
|
|