working code

This commit is contained in:
matan1905
2023-11-02 18:12:17 +02:00
commit 3aa7037fc7
6 changed files with 255 additions and 0 deletions
+1
View File
@@ -0,0 +1 @@
__pycache__/*
+6
View File
@@ -0,0 +1,6 @@
import __main__
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
+16
View File
@@ -0,0 +1,16 @@
import discord
from discord.ext import commands
intents = discord.Intents.default()
intents.message_content = True
discord_client = commands.Bot(command_prefix='!', intents=intents)
# Event handler for when the bot is ready
@discord_client.event
async def on_ready():
print(f'Logged in as {discord_client.user.name}. Ready to take requests!')
+199
View File
@@ -0,0 +1,199 @@
import time
from .serving_manager import serving_manager
import threading
from .discord_client import discord_client
import threading
from collections import deque
from .utils import parse_command_string, tensorToImageConversion
import discord
import io
import asyncio
class ServingOutput:
def __init__(self):
# start listening to api/discord
# when something happen, pass to serving manager with the details
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"serving_config": ("SERVING_CONFIG",),
"image": ("IMAGE",),
},
}
RETURN_TYPES = ()
# RETURN_NAMES = ("image_output_name",)
FUNCTION = "out"
OUTPUT_NODE = True
CATEGORY = "Serving-Toolkit"
def out(self, image,serving_config):
serving_config["serve_image_function"](image)
return {}
class ServingInputText:
def __init__(self):
# start listening to api/discord
# when something happen, pass to serving manager with the details
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"serving_config": ("SERVING_CONFIG",),
"argument": ("STRING", {
"multiline": False,
"default": "prompt"
}),
"default": ("STRING", {
"multiline": True,
"default": ""
}),
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("text",)
FUNCTION = "out"
CATEGORY = "Serving-Toolkit"
def out(self, serving_config, argument,default):
if argument not in serving_config:
return (default,)
return (serving_config[argument],)
class ServingInputNumber:
def __init__(self):
# start listening to api/discord
# when something happen, pass to serving manager with the details
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"serving_config": ("SERVING_CONFIG",),
"argument": ("STRING", {
"multiline": False,
"default": "number"
}),
"default": ("FLOAT", {"default": 0.0,}),
"min_value": ("FLOAT", {"default": -999999.0,}),
"max_value": ("FLOAT", {"default": 999999.0,}),
"step": ("FLOAT", {"default": 0.1,}),
}
}
RETURN_TYPES = ("FLOAT", "INT")
FUNCTION = "out"
CATEGORY = "Serving-Toolkit"
def out(self, serving_config, argument,default, min_value, max_value, step):
val = default
if argument in serving_config and serving_config[argument].replace('.','',1).isdigit():
val = serving_config[argument]
valFloat = min(max(float(val), min_value), max_value) // step * step
valInt = round(valFloat)
return (valFloat,valInt)
class DiscordServing():
discord_running = False
def __init__(self):
self.registered_command = False
self.data_ready = threading.Event()
self.data = deque()
self.discord_token = None
pass
def discord_runner(self):
discord_client.run(self.discord_token)
def get_data(self):
if not self.data:
self.data_ready.wait()
data = self.data.popleft()
self.data_ready.clear()
return data
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"discord_token": ("STRING", {
"multiline": True,
"default": ""
}),
"command_name": ("STRING", {
"multiline": False,
"default": "generate"
})
}
}
RETURN_TYPES = ("SERVING_CONFIG",)
RETURN_NAMES = ("Serving config",)
FUNCTION = "serve"
@classmethod
def IS_CHANGED(cls, **kwargs):
return float("NaN")
# OUTPUT_NODE = False
CATEGORY = "Serving-Toolkit"
def serve(self, command_name, discord_token):
if not DiscordServing.discord_running:
self.discord_token = discord_token
run_discord = threading.Thread(target=self.discord_runner)
run_discord.start()
print("Client running")
DiscordServing.discord_running = True
if not self.registered_command:
self.registered_command = True
@discord_client.command(name=command_name)
async def execute(ctx):
parsed_data = parse_command_string(ctx.message.content,command_name)
def serve_image_function(image):
image_file = tensorToImageConversion(image)
asyncio.run_coroutine_threadsafe(ctx.reply(file=discord.File(image_file, filename='image.png')), discord_client.loop)
parsed_data["serve_image_function"] = serve_image_function
self.data.append(parsed_data)
self.data_ready.set()
data = self.get_data()
return (data,)
# A dictionary that contains all nodes you want to export with their names
# NOTE: names should be globally unique
NODE_CLASS_MAPPINGS = {
"ServingOutput": ServingOutput,
"ServingInputText": ServingInputText,
"ServingInputNumber": ServingInputNumber,
"DiscordServing": DiscordServing
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
"ServingOutput": "Serving Output",
"DiscordServing": "Discord Serving",
"ServingInputText": "Serving Input Text",
"ServingInputNumber": "Serving Input Number",
}
# input - simply a push
+1
View File
@@ -0,0 +1 @@
discord.py
+32
View File
@@ -0,0 +1,32 @@
def parse_command_string(command_string, command_name):
textAndArgs = command_string[1+ len(command_name):].strip().split('--')
result = {}
text = textAndArgs[0].strip()
args = textAndArgs[1:]
print(args)
# The first element is the "freeText" part, remove any leading or trailing whitespace.
result["prompt"] = text.strip()
for arg in args:
parts = arg.split()
if len(parts) > 1:
# Extract the argument name and value
arg_name = parts[0].strip()
arg_value = ' '.join(parts[1:]).strip()
result[arg_name] = arg_value
return result
import io
from PIL import Image
def tensorToImageConversion(img_tensor ):
img_tensor = (img_tensor * 255).byte()
img_array = img_tensor.squeeze(0).numpy()
img_pil = Image.fromarray(img_array)
img_byte_array = io.BytesIO()
img_pil.save(img_byte_array, format='PNG')
img_byte_array.seek(0)
return img_byte_array