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