ComfyUI/memedeck-v1.py

import asyncio
import base64
from io import BytesIO
import os
import logging
import signal
import struct
from typing import Optional
import uuid
from PIL import Image, ImageOps
from functools import partial

import pika
import json

import requests
import aiohttp

# load from env file
# load from .env file
from dotenv import load_dotenv
load_dotenv()

amqp_addr = os.getenv('AMQP_ADDR') or 'amqp://api:gacdownatravKekmy9@51.8.120.154:5672/dev'

# define the enum in python
from enum import Enum

class QueueProgressKind(Enum):
    # make json serializable
    ImageGenerated = "image_generated"
    ImageGenerating = "image_generating"
    SamePrompt = "same_prompt"
    FaceswapGenerated = "faceswap_generated"
    FaceswapGenerating = "faceswap_generating"
    Failed = "failed"

class MemedeckWorker:
    class BinaryEventTypes:
        PREVIEW_IMAGE = 1
        UNENCODED_PREVIEW_IMAGE = 2
        
    class JsonEventTypes(Enum):
        PROGRESS = "progress"
        EXECUTING = "executing"
        EXECUTED = "executed"
        ERROR = "error"
        STATUS = "status"
        
    """
    MemedeckWorker is a class that is responsible for relaying messages between comfy and the memedeck backend api
    it is used to send images to the memedeck backend api and to receive prompts from the memedeck backend api
    """
    def __init__(self, loop):
        MemedeckWorker.instance = self
        # set logging level to info 
        logging.getLogger().setLevel(logging.INFO)
        self.active_tasks_map = {}
        self.current_task = None
        
        self.client_id = None
        self.ws_id = None
        self.websocket_node_id = None
        self.current_node = None
        self.current_progress = 0
        self.current_context = None
        
        self.loop = loop
        self.messages = asyncio.Queue()

        self.http_client = None
        self.prompt_queue = None
        self.validate_prompt = None
        self.last_prompt_id = None
        
        self.amqp_url = amqp_addr
        self.queue_name = os.getenv('QUEUE_NAME') or 'generic-queue'
        self.api_url = os.getenv('API_ADDRESS') or 'http://0.0.0.0:8079/v2'
        self.api_key = os.getenv('API_KEY') or 'eb46e20a-cc25-4ed4-a39b-f47ca8ff3383'
        
        # Configure logging
        logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
        self.logger = logging.getLogger(__name__)
        self.logger.info(f"\n[memedeck]: initialized with API URL: {self.api_url} and API Key: {self.api_key}\n")
        
    def on_connection_open(self, connection):
        self.connection = connection
        self.connection.channel(on_open_callback=self.on_channel_open)

    def on_channel_open(self, channel):
        self.channel = channel
        
        # only consume one message at a time
        self.channel.basic_qos(prefetch_size=0, prefetch_count=1)
        self.channel.queue_declare(queue=self.queue_name, durable=True)
        self.channel.basic_consume(queue=self.queue_name, on_message_callback=self.on_message_received)
        
    def start(self, prompt_queue, validate_prompt):
      self.prompt_queue = prompt_queue
      self.validate_prompt = validate_prompt
      
      parameters = pika.URLParameters(self.amqp_url)
      logging.getLogger('pika').setLevel(logging.WARNING) # supress all logs from pika
      self.connection = pika.SelectConnection(parameters, on_open_callback=self.on_connection_open)

      try:
          self.connection.ioloop.start()
      except KeyboardInterrupt:
          self.connection.close()
          self.connection.ioloop.start()
      
    def on_message_received(self, channel, method, properties, body):        
        decoded_string = body.decode('utf-8')
        json_object = json.loads(decoded_string)
        payload = json_object[1]
              
        # execute the task
        prompt = payload["nodes"]
        valid =  self.validate_prompt(prompt)
        
        self.current_node = None
        self.current_progress = 0
        self.websocket_node_id = None
        self.ws_id = payload["source_ws_id"]
        self.current_context = payload["req_ctx"]
        
        for node in prompt: # search through prompt nodes for websocket_node_id
            if isinstance(prompt[node], dict) and prompt[node].get("class_type") == "SaveImageWebsocket":
                self.websocket_node_id = node
                break
            
        if valid[0]:
          prompt_id = str(uuid.uuid4())
          outputs_to_execute = valid[2]
          self.active_tasks_map[payload["source_ws_id"]] = {
            "prompt_id": prompt_id,
            "prompt": prompt,
            "outputs_to_execute": outputs_to_execute,
            "client_id": "memedeck-1",
            "is_memedeck": True,
            "websocket_node_id": self.websocket_node_id,
            "ws_id": payload["source_ws_id"],
            "context": payload["req_ctx"],
            "current_node": None,
            "current_progress": 0,
          }
          self.prompt_queue.put((0, prompt_id, prompt, { 
                "client_id": "memedeck-1", 
                'is_memedeck': True, 
                'websocket_node_id': self.websocket_node_id, 
                'ws_id': payload["source_ws_id"], 
                'context': payload["req_ctx"] 
            }, outputs_to_execute))
          self.set_last_prompt_id(prompt_id)
          channel.basic_ack(delivery_tag=method.delivery_tag) # ack the task        
        else:
          channel.basic_nack(delivery_tag=method.delivery_tag, requeue=False) # unack the message
    
    # --------------------------------------------------
    # callbacks for the prompt queue
    # --------------------------------------------------
    def queue_updated(self):
        # print json of the queue info but only print the first 100 lines
        info = self.get_queue_info()
        # update_type = info['']
        # self.send_sync("status", { "status": self.get_queue_info() })
          
    def get_queue_info(self):
        prompt_info = {}
        exec_info = {}
        exec_info['queue_remaining'] = self.prompt_queue.get_tasks_remaining()
        prompt_info['exec_info'] = exec_info
        return prompt_info
    
    def send_sync(self, event, data, sid=None):
        
        self.loop.call_soon_threadsafe(
            self.messages.put_nowait, (event, data, sid))
        
    def set_last_prompt_id(self, prompt_id):
        self.last_prompt_id = prompt_id
          
    async def publish_loop(self):
        while True:
            msg = await self.messages.get()
            await self.send(*msg)
            
    async def send(self, event, data, sid=None):
        current_task = self.active_tasks_map.get(sid)
        if current_task is None or current_task['ws_id'] != sid:
            return
        
        if event == MemedeckWorker.BinaryEventTypes.UNENCODED_PREVIEW_IMAGE: # preview and unencoded images are sent here
            self.logger.info(f"[memedeck]: sending image preview for {sid}")
            await self.send_preview(data, sid=current_task['ws_id'], progress=current_task['current_progress'], context=current_task['context'])
        else: # send json data / text data
            if event == "executing":
                current_task['current_node'] = data['node']   
            elif event == "executed":
                self.logger.info(f"---> [memedeck]: executed event for {sid}")
                prompt_id = data['prompt_id']
                if prompt_id in self.active_tasks_map:
                    del self.active_tasks_map[prompt_id]
            elif event == "progress":
                if current_task['current_node'] == current_task['websocket_node_id']: # if the node is the websocket node, then set the progress to 100
                    current_task['current_progress'] = 100
                else: # if the node is not the websocket node, then set the progress to the progress from the node
                    current_task['current_progress'] = data['value'] / data['max'] * 100
                    if current_task['current_progress'] == 100 and current_task['current_node'] != current_task['websocket_node_id']:
                        # in case the progress is 100 but the node is not the websocket node, then set the progress to 95
                        current_task['current_progress'] = 95 # this allows the full resolution image to be sent on the 100 progress event
                
                if data['value'] == 1: # if the value is 1, then send started to api
                    start_data = {
                        "ws_id": current_task['ws_id'],
                        "status": "started",
                        "info": None,
                    }
                    await self.send_to_api(start_data)
                    
            elif event == "status":
                self.logger.info(f"[memedeck]: sending status event: {data}")
            
        self.active_tasks_map[sid] = current_task

        
    async def send_preview(self, image_data, sid=None, progress=None, context=None):
        # if self.current_progress is odd, then don't send the preview
        if progress % 2 == 1:
            return
                
        image_type = image_data[0]
        image = image_data[1]
        max_size = image_data[2]
        if max_size is not None:
            if hasattr(Image, 'Resampling'):
                resampling = Image.Resampling.BILINEAR
            else:
                resampling = Image.ANTIALIAS

            image = ImageOps.contain(image, (max_size, max_size), resampling)

        bytesIO = BytesIO()
        image.save(bytesIO, format=image_type, quality=100 if progress == 96 else 75, compress_level=1)
        preview_bytes = bytesIO.getvalue()
            
        ai_queue_progress = {
            "ws_id": sid,
            "kind": "image_generating" if progress < 100 else "image_generated",
            "data": list(preview_bytes),
            "progress": int(progress),
            "context": context
        }
       
        await self.send_to_api(ai_queue_progress)
        
    async def send_to_api(self, data):
        if self.websocket_node_id is None: # check if the node is still running
            logging.error(f"[memedeck]: websocket_node_id is None for {data['ws_id']}")
            return
         
        try:
            post_func = partial(requests.post, f"{self.api_url}/generation/update", json=data)        
            await self.loop.run_in_executor(None, post_func)
        except Exception as e:
            self.logger.error(f"[memedeck]: error sending to api: {e}")