{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import websocket \n",
    "import uuid\n",
    "import json\n",
    "import urllib.request\n",
    "import urllib.parse\n",
    "\n",
    "import json\n",
    "\n",
    "import glob\n",
    "import os\n",
    "\n",
    "class LP:\n",
    "    def __init__(self , listen = '127.0.0.1' , port = '8188' , work_path = './workflow/lp.json'):\n",
    "\n",
    "        with open(work_path, \"r\", encoding=\"utf-8\") as f:\n",
    "            workflow_data = f.read()\n",
    "        self.lp_workflow = json.loads(workflow_data)\n",
    "        self.server_address = f'{listen}:{port}'\n",
    "\n",
    "    def queue_prompt(self , prompt , client_id):\n",
    "        p = {\"prompt\": prompt, \"client_id\": client_id}\n",
    "        data = json.dumps(p).encode('utf-8')\n",
    "        req =  urllib.request.Request(\"http://{}/prompt\".format(self.server_address), data=data)\n",
    "        return json.loads(urllib.request.urlopen(req).read())\n",
    "\n",
    "    def get_history(self , prompt_id):\n",
    "        with urllib.request.urlopen(\"http://{}/history/{}\".format(self.server_address, prompt_id)) as response:\n",
    "            return json.loads(response.read())\n",
    "    \n",
    "\n",
    "    def LP(self, image_path, video_path):\n",
    "        try:\n",
    "            # Create a copy of the current workflow and generate a unique client_id\n",
    "            current_workflow = self.lp_workflow.copy()\n",
    "            client_id = str(uuid.uuid4())\n",
    "\n",
    "            # Update the current workflow with the provided image and video paths\n",
    "            current_workflow[\"168\"]['inputs']['filename_prefix'] = client_id\n",
    "            current_workflow['208']['inputs']['image'] = image_path\n",
    "            current_workflow['202']['inputs']['video'] = video_path\n",
    "\n",
    "            # Try connecting to the WebSocket server\n",
    "            try:\n",
    "                ws = websocket.WebSocket()\n",
    "                ws.connect(\"ws://{}/ws?clientId={}\".format(self.server_address, client_id))\n",
    "            except Exception as e:\n",
    "                print(f\"Error connecting to WebSocket server: {e}\")\n",
    "                return None\n",
    "\n",
    "            # Try getting the output from the WebSocket\n",
    "            try:\n",
    "                video_output = self.get_output(ws, current_workflow, client_id)\n",
    "            except Exception as e:\n",
    "                print(f\"Error during execution of workflow: {e}\")\n",
    "                return None\n",
    "\n",
    "            # Search for files generated with the client_id prefix\n",
    "            try:\n",
    "                folder_path = '/home/itek/Desktop/ComfyUI/output'\n",
    "                beginning_string = client_id\n",
    "                paths = glob.glob(os.path.join(folder_path, beginning_string + '*'))\n",
    "            except Exception as e:\n",
    "                print(f\"Error searching for output files: {e}\")\n",
    "                return None\n",
    "\n",
    "            # Return the video output and paths if everything was successful\n",
    "            return video_output, paths\n",
    "\n",
    "        except Exception as e:\n",
    "            # Catch any unforeseen errors\n",
    "            print(f\"An unexpected error occurred: {e}\")\n",
    "            return None\n",
    "\n",
    "    def get_output(self, ws, prompt, client_id):\n",
    "        try:\n",
    "            # Send the prompt and wait for the response\n",
    "            prompt_id = self.queue_prompt(prompt, client_id)['prompt_id']\n",
    "            while True:\n",
    "                out = ws.recv()\n",
    "                if isinstance(out, str):\n",
    "                    message = json.loads(out)\n",
    "                    if message['type'] == 'executing':\n",
    "                        data = message['data']\n",
    "                        if data['node'] is None and data['prompt_id'] == prompt_id:\n",
    "                            break\n",
    "                else:\n",
    "                    continue\n",
    "\n",
    "            # Get the output file from the history\n",
    "            history = self.get_history(prompt_id)[prompt_id]\n",
    "            output_file = history['outputs']['168']['gifs'][0]['filename']\n",
    "\n",
    "            return output_file\n",
    "\n",
    "        except Exception as e:\n",
    "            # Handle errors during WebSocket communication or workflow processing\n",
    "            print(f\"Error in get_output function: {e}\")\n",
    "            return None\n",
    "\n",
    "    # def get_output(self , ws , prompt , client_id):\n",
    "    #     prompt_id = self.queue_prompt(prompt , client_id)['prompt_id']\n",
    "    #     while True:\n",
    "    #         out = ws.recv()\n",
    "    #         if isinstance(out, str):\n",
    "    #             message = json.loads(out)\n",
    "    #             if message['type'] == 'executing':\n",
    "    #                 data = message['data']\n",
    "    #                 if data['node'] is None and data['prompt_id'] == prompt_id:\n",
    "    #                     break\n",
    "    #         else:\n",
    "    #             continue\n",
    "\n",
    "    #     history = self.get_history(prompt_id)[prompt_id]\n",
    "    #     output_file = history['outputs']['168']['gifs'][0]['filename']\n",
    "\n",
    "    #     return output_file\n",
    "\n",
    "    # def LP(self , image_path , video_path):\n",
    "\n",
    "    #     current_workflow = self.lp_workflow.copy()\n",
    "    #     client_id = str(uuid.uuid4())\n",
    "\n",
    "    #     current_workflow[\"168\"]['inputs']['filename_prefix'] = client_id\n",
    "    #     current_workflow['208']['inputs']['image'] = image_path\n",
    "    #     current_workflow['202']['inputs']['video'] = video_path\n",
    "\n",
    "    #     ws = websocket.WebSocket()\n",
    "    #     ws.connect(\"ws://{}/ws?clientId={}\".format(self.server_address, client_id))\n",
    "    #     video_output = self.get_output(ws, current_workflow , client_id)\n",
    "        \n",
    "    #     folder_path = '/home/itek/Desktop/ComfyUI/output'\n",
    "    #     beginning_string = client_id\n",
    "    #     paths = glob.glob(os.path.join(folder_path, beginning_string + '*'))\n",
    "\n",
    "\n",
    "    #     return video_output , paths\n",
    "\n",
    "    # def __del__(self):\n",
    "    #     self.model = None\n",
    "    #     self.processor = None\n",
    "    #     del self.model\n",
    "    #     del self.processor\n",
    "    #     torch.cuda.empty_cache()\n",
    "    #     import gc\n",
    "    #     gc.collect()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import socket\n",
    "import time\n",
    "import requests\n",
    "import subprocess\n",
    "import os\n",
    "\n",
    "class ComfyUI:\n",
    "    def __init__(self):\n",
    "        self.process = None\n",
    "        self.port = None\n",
    "\n",
    "    def is_port_in_use(self , port , host='127.0.0.1'):\n",
    "        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:\n",
    "            s.settimeout(1)\n",
    "            try:\n",
    "                s.bind((host, port))\n",
    "            except socket.error:\n",
    "                return True  \n",
    "            return False  \n",
    "    \n",
    "    def is_server_ready(self , url, timeout=30):\n",
    "        start_time = time.time()\n",
    "        while True:\n",
    "            try:\n",
    "                response = requests.get(url)\n",
    "                if response.status_code == 200:\n",
    "                    return True\n",
    "            except requests.RequestException:\n",
    "                pass\n",
    "            \n",
    "            if time.time() - start_time > timeout:\n",
    "                return False\n",
    "            time.sleep(1) \n",
    "    \n",
    "    def SS(self):\n",
    "        try:\n",
    "            port_to_check = 60000\n",
    "            con_count = 0\n",
    "            while(True):\n",
    "                if con_count > 100:\n",
    "                    return False\n",
    "                if self.is_port_in_use(port_to_check):\n",
    "                    port_to_check = port_to_check + 1\n",
    "                    con_count = con_count + 1\n",
    "                else:\n",
    "                    break\n",
    "                \n",
    "            port_to_check = str(port_to_check)\n",
    "            script_dir = \"./ComfyUI/ComfyUI\"\n",
    "            script_dir = os.path.abspath(script_dir)\n",
    "\n",
    "            code_dir = \"./ComfyUI/ComfyUI/main.py\"\n",
    "            code_dir = os.path.abspath(code_dir)\n",
    "\n",
    "            command = [\"/home/itek/miniconda3/envs/ai/bin/python\", os.path.abspath(code_dir), \"--port\", port_to_check, '--listen' , '127.0.0.1']\n",
    "            self.process = subprocess.Popen(command, cwd=script_dir , stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)\n",
    "            self.port = port_to_check\n",
    "\n",
    "            return True\n",
    "                \n",
    "        except Exception as e:  \n",
    "            print(e)\n",
    "            return False\n",
    "        \n",
    "    def restart_if_down(self):\n",
    "\n",
    "        if self.port is None:\n",
    "            return False\n",
    "        \n",
    "        server_url = f\"http://127.0.0.1:{self.port}\"\n",
    "        if not self.is_server_ready(server_url, timeout=5):\n",
    "            \n",
    "            if self.process is not None:\n",
    "                self.process.terminate()\n",
    "                self.process.kill()\n",
    "                self.process = None\n",
    "                self.port = None\n",
    "\n",
    "            return self.SS()\n",
    "\n",
    "        return True\n",
    "\n",
    "    def __del__(self):\n",
    "        if self.process is not None:\n",
    "            self.process.terminate()\n",
    "            self.process.kill()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "ui = ComfyUI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Errno 2] No such file or directory: '/home/itek/Desktop/api comfyui/scripts/ComfyUI/ComfyUI'\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ui.SS()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "lp = LP()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error connecting to WebSocket server: [Errno 111] Connection refused\n"
     ]
    }
   ],
   "source": [
    "lp.LP('/home/itek/Desktop/api comfyui/121.png' , '/home/itek/Desktop/api comfyui/d0.mp4')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import socket\n",
    "import subprocess\n",
    "import requests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "port_to_check = 60000\n",
    "con_count = 0\n",
    "while(True):\n",
    "    if con_count > 100:\n",
    "        break\n",
    "    if is_port_in_use(port_to_check) or (search_port(port_to_check)):\n",
    "        port_to_check = port_to_check + 1\n",
    "        con_count = con_count + 1\n",
    "    else:\n",
    "        break\n",
    "\n",
    "    port_to_check = str(port_to_check)\n",
    "    script_dir = \"./subprocess/LivePortrait\"\n",
    "    script_dir = os.path.abspath(script_dir)\n",
    "\n",
    "    command = [\"/home/itek/miniconda3/envs/ai/bin/python\", script_dir+\"/app.py\", \"--server_port\", port_to_check]\n",
    "    process = subprocess.Popen(command, cwd=script_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "command = [\"/home/itek/miniconda3/envs/ai/bin/python\", \"/home/itek/Desktop/api comfyui/ComfyUI/ComfyUI/main.py\", \"--port\", '8000' , '--listen' , '127.0.0.1']\n",
    "process = subprocess.Popen(command, cwd='/home/itek/Desktop/api comfyui/ComfyUI/ComfyUI' , stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/home/itek/Desktop/api comfyui/scripts/ComfyUI'"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "script_dir = \"./ComfyUI\"\n",
    "os.path.abspath(script_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "process.terminate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_port_lp():\n",
    "    try:\n",
    "        if not get_model_port_lp():\n",
    "            port_to_check = 60000\n",
    "            con_count = 0\n",
    "            while(True):\n",
    "                if con_count > 100:\n",
    "                    break\n",
    "                if is_port_in_use(port_to_check) or (search_port(port_to_check)):\n",
    "                    port_to_check = port_to_check + 1\n",
    "                    con_count = con_count + 1\n",
    "                else:\n",
    "                    break\n",
    "            port_to_check = str(port_to_check)\n",
    "            script_dir = \"./subprocess/LivePortrait\"\n",
    "            script_dir = os.path.abspath(script_dir)\n",
    "\n",
    "            command = [\"/home/itek/miniconda3/envs/ai/bin/python\", script_dir+\"/app.py\", \"--server_port\", port_to_check]\n",
    "            process = subprocess.Popen(command, cwd=script_dir)\n",
    "\n",
    "            global model_port_lp\n",
    "            model_port_lp = {\"port\":port_to_check , \"pid\":process}\n",
    "\n",
    "            insert_port(port_to_check)\n",
    "\n",
    "            is_server_ready(f'http://127.0.0.1:{port_to_check}/')\n",
    "\n",
    "\n",
    "        return True\n",
    "    except:\n",
    "        return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "listen = '127.0.0.1'\n",
    "port = '8188'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "is_server_ready('http://127.0.0.1:8188')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#model lp \n",
    "import time\n",
    "import socket\n",
    "import subprocess\n",
    "import requests\n",
    "\n",
    "\n",
    "def is_port_in_use(port, host='localhost'):\n",
    "    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:\n",
    "        s.settimeout(1)\n",
    "        try:\n",
    "            s.bind((host, port))\n",
    "        except socket.error:\n",
    "            return True  \n",
    "        return False  \n",
    "    \n",
    "def is_server_ready(url, timeout=30):\n",
    "    start_time = time.time()\n",
    "    while True:\n",
    "        try:\n",
    "            response = requests.get(url)\n",
    "            if response.status_code == 200:\n",
    "                return True\n",
    "        except requests.RequestException:\n",
    "            pass\n",
    "        \n",
    "        if time.time() - start_time > timeout:\n",
    "            return False\n",
    "        time.sleep(1) \n",
    "\n",
    "model_port_lp = None\n",
    "\n",
    "def load_port_lp():\n",
    "    try:\n",
    "        if not get_model_port_lp():\n",
    "            port_to_check = 60000\n",
    "            con_count = 0\n",
    "            while(True):\n",
    "                if con_count > 100:\n",
    "                    break\n",
    "                if is_port_in_use(port_to_check) or (search_port(port_to_check)):\n",
    "                    port_to_check = port_to_check + 1\n",
    "                    con_count = con_count + 1\n",
    "                else:\n",
    "                    break\n",
    "            port_to_check = str(port_to_check)\n",
    "            script_dir = \"./subprocess/LivePortrait\"\n",
    "            script_dir = os.path.abspath(script_dir)\n",
    "\n",
    "            command = [\"/home/itek/miniconda3/envs/ai/bin/python\", script_dir+\"/app.py\", \"--server_port\", port_to_check]\n",
    "            process = subprocess.Popen(command, cwd=script_dir)\n",
    "\n",
    "            global model_port_lp\n",
    "            model_port_lp = {\"port\":port_to_check , \"pid\":process}\n",
    "\n",
    "            insert_port(port_to_check)\n",
    "\n",
    "            is_server_ready(f'http://127.0.0.1:{port_to_check}/')\n",
    "\n",
    "\n",
    "        return True\n",
    "    except:\n",
    "        return False\n",
    "\n",
    "def unload_port_lp():\n",
    "    try:\n",
    "        if get_model_port_lp():\n",
    "            global model_port_lp\n",
    "            port = model_port_lp['port']\n",
    "            remove_port(port)\n",
    "            pid = model_port_lp['pid']\n",
    "            pid.terminate()\n",
    "            pid.kill()\n",
    "            model_port_lp = None\n",
    "            return True\n",
    "        else:\n",
    "            return True\n",
    "    except:\n",
    "        return False\n",
    "    \n",
    "def get_model_port_lp():\n",
    "    global model_port_lp\n",
    "    if model_port_lp is None:\n",
    "        return False\n",
    "    return True\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}