ig_bot/SD.py

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