Initial

.env

@@ -0,0 +1,8 @@
+SERVICE_ENDPOINT=http://0.0.0.0/api/
+
+MODEL_NAME=runwayml/stable-diffusion-v1-5
+
+APP_DIR=/Users/akin/Desktop/txt2img-consumer/
+INSTANCE_DIR=/Users/akin/Desktop/txt2img-consumer/instance/
+OUTPUT_DIR=/content/model
+CLASS_DIR=/content/data/person

.gitignore

@@ -0,0 +1,4 @@
+.idea
+instance
+instance.zip
+__pycache__

Service.py

@@ -0,0 +1,51 @@
+import requests
+import os, shutil
+from dotenv import load_dotenv
+import zipfile
+
+
+class Service:
+    def __init__(self, orderId, hash):
+        load_dotenv()
+
+        self.service_endpoint = os.getenv('SERVICE_ENDPOINT')
+        self.app_dir = os.getenv('APP_DIR')
+        self.instance_dir = os.getenv('INSTANCE_DIR')
+        self.archive_filename = self.app_dir + 'instance.zip'
+
+        self.headers = {
+            'Content-Type': 'application/json',
+            'Authorization': hash,
+        }
+
+        self.params = {
+            'orderId': orderId,
+        }
+
+    # Download images
+    def download(self):
+        self.scarify()
+
+        # Get images
+        response = requests.post(self.service_endpoint + 'download', headers=self.headers, json=self.params, verify=False)
+        with open(self.archive_filename, 'wb') as f:
+            f.write(response.content)
+
+        # Extract
+        with zipfile.ZipFile(self.archive_filename, "r") as zip_ref:
+            zip_ref.extractall(self.instance_dir)
+
+        print("Downloaded images")
+
+    # Get data
+    def data(self):
+        response = requests.post(self.service_endpoint + 'data', headers=self.headers, json=self.params, verify=False)
+        return response.json()
+
+    # Remove old files
+    def scarify(self):
+        if os.path.exists(self.instance_dir):
+            shutil.rmtree(self.instance_dir)
+
+        if os.path.exists(self.archive_filename):
+            os.remove(self.archive_filename)

Training.py

@@ -0,0 +1,24 @@
+import subprocess, os
+
+
+class Training:
+    def __init__(self):
+        self.app_path = os.getenv('APP_DIR')
+        self.model_name = os.getenv('MODEL_NAME')
+        self.instance_dir = os.getenv('INSTANCE_DIR')
+        self.output_dir = os.getenv('OUTPUT_DIR')
+        self.class_dir = os.getenv('CLASS_DIR')
+
+        self.start()
+
+    def start(self):
+        args = [
+            self.app_path + 'training.sh',
+            self.model_name,
+            self.instance_dir,
+            self.output_dir,
+            self.class_dir
+        ]
+        subprocess.run(['chmod', '+x', self.app_path + 'training.sh'], cwd=self.app_path)
+        process = subprocess.Popen(args, cwd=self.app_path)
+        process.wait()

handler.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python
+
+import pika
+import json
+import phpserialize
+from Service import Service
+from Training import Training
+
+
+class EndpointHandler:
+    def __init__(self):
+        credentials = pika.PlainCredentials('admin', '123456')
+        self.connection = pika.BlockingConnection(pika.ConnectionParameters('0.0.0.0', 5672, '/', credentials))
+        self.channel = self.connection.channel()
+        self.channel.queue_declare(queue='avator', durable=True)
+        self.consume()
+
+    def consume(self):
+        print(' [*] Waiting for messages. To exit press CTRL+C')
+
+        def callback(ch, method, properties, body):
+            payload = json.loads(body)
+            service = Service(orderId=payload['orderId'], hash=payload['hash'])
+
+            # Download images
+            service.download()
+
+            # Get data
+            data = service.data()
+
+            # Start training
+            Training()
+
+            # ch.basic_ack(delivery_tag=method.delivery_tag)
+
+        self.channel.basic_qos(prefetch_count=1)
+        self.channel.basic_consume(queue='avator', on_message_callback=callback)
+        self.channel.start_consuming()
+
+
+worker = EndpointHandler()

requirements.txt

@@ -0,0 +1,9 @@
+pika
+phpserialize
+torchvision
+tensorboard
+modelcards
+wget
+bitsandbytes
+xformers
+natsort

setup.py

@@ -0,0 +1,56 @@
+from distutils.command.build import build as _build
+import setuptools
+import subprocess
+
+print("123")
+raise None
+
+
+CUSTOM_COMMANDS = [
+    #['apt-get', 'install', 'wget'],
+    #['python', '-c', "from huggingface_hub.hf_api import HfFolder; HfFolder.save_token('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')"],
+    ['wget', '-q', "https://github.com/ShivamShrirao/diffusers/raw/main/examples/dreambooth/train_dreambooth.py"],
+    ['wget', '-q', "https://github.com/ShivamShrirao/diffusers/raw/main/scripts/convert_diffusers_to_original_stable_diffusion.py"],
+]
+
+
+class Build(_build):
+    sub_commands = _build.sub_commands + [('CustomCommands', None)]
+
+
+class CustomCommands(setuptools.Command):
+    def initialize_options(self):
+        pass
+
+    def finalize_options(self):
+        pass
+
+    def RunCustomCommand(self, command_list):
+        print('Running command: %s' % command_list)
+        p = subprocess.Popen(
+            command_list,
+            stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
+        )
+
+        stdout_data, _ = p.communicate()
+        print( 'Command output: %s' % stdout_data)
+
+        if p.returncode != 0:
+            raise RuntimeError(
+                'Command %s failed: exit code: %s' % (command_list, p.returncode))
+
+    def run(self):
+        for command in CUSTOM_COMMANDS:
+            self.RunCustomCommand(command)
+
+
+setup(
+    name="Txt2ImgConsumer",
+    version="1.0",
+    description="",
+    packages=find_packages(),
+    cmdclass={
+        'build': Build,
+        'CustomCommands': CustomCommands,
+    }
+)

training.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+
+model_name=$1
+instance_dir=$2
+output_dir=$3
+class_dir=$4