#!/usr/bin/env python
# -*- coding: utf-8 -*-

# This script relies on large code sections located at:
# https://github.com/comfyanonymous/ComfyUI/blob/master/script_examples/websockets_api_example.py

# Other code sections are (C) Copyright 2023, P. Lutus: https://www.arachnoid.com
# and are released under the GPL: https://www.gnu.org/licenses/gpl-3.0.en.html

# This is the ComfyUI version of the matrix script

# must pip install websocket-client
import websocket
import uuid
import json
import urllib.request
import urllib.parse
import random
import io
import time
from PIL import Image
from PIL.PngImagePlugin import PngInfo
from dataclasses import dataclass

class MultiDraw:
  def __init__(self):
    
    self.server='localhost:8188'
    self.client_id = str(uuid.uuid4())
    self.savepath='image_matrix'
    
    # acquire this with save (api format) button
    # to make this visible click gear icon then "enable dev mode options"
    
    json_text = """
    {
        "3": {
            "class_type": "KSampler",
            "inputs": {
                "cfg": 8,
                "denoise": 1,
                "latent_image": [
                    "5",
                    0
                ],
                "model": [
                    "4",
                    0
                ],
                "negative": [
                    "7",
                    0
                ],
                "positive": [
                    "6",
                    0
                ],
                "sampler_name": "euler",
                "scheduler": "normal",
                "seed": 8566257,
                "steps": 20
            }
        },
        "4": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {
                "ckpt_name": "v1-5-pruned-emaonly.ckpt"
            }
        },
        "5": {
            "class_type": "EmptyLatentImage",
            "inputs": {
                "batch_size": 1,
                "height": 1024,
                "width": 1024
            }
        },
        "6": {
            "class_type": "CLIPTextEncode",
            "inputs": {
                "clip": [
                    "4",
                    1
                ],
                "text": "masterpiece best quality girl"
            }
        },
        "7": {
            "class_type": "CLIPTextEncode",
            "inputs": {
                "clip": [
                    "4",
                    1
                ],
                "text": "bad hands"
            }
        },
        "8": {
            "class_type": "VAEDecode",
            "inputs": {
                "samples": [
                    "3",
                    0
                ],
                "vae": [
                    "4",
                    2
                ]
            }
        },
        "9": {
            "class_type": "SaveImage",
            "inputs": {
                "filename_prefix": "ComfyUI",
                "images": [
                    "8",
                    0
                ]
            }
        }
    }
    """
    self.json_prompt = json.loads(json_text)
    
  def queue_prompt(self,prompt):
    p = {"prompt": prompt, "client_id": self.client_id}
    data = json.dumps(p).encode('utf-8')
    req =  urllib.request.Request(f"http://{self.server}/prompt", data=data)
    return json.loads(urllib.request.urlopen(req).read())

  def get_image(self,filename, subfolder, folder_type):
    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
    url_values = urllib.parse.urlencode(data)
    with urllib.request.urlopen(f"http://{self.server}/view?{url_values}") as response:
        return response.read()

  def get_history(self,prompt_id):
    with urllib.request.urlopen(f"http://{self.server}/history/{prompt_id}") as response:
        return json.loads(response.read())

  def get_images(self,ws, prompt):
    prompt_id = self.queue_prompt(prompt)['prompt_id']
    output_images = {}
    while True:
        out = ws.recv()
        if isinstance(out, str):
            message = json.loads(out)
            if message['type'] == 'executing':
                data = message['data']
                if data['node'] is None and data['prompt_id'] == prompt_id:
                    break 
        else:
            continue 

    history = self.get_history(prompt_id)[prompt_id]
    prompt = history['prompt'][2]
    metadata = PngInfo()
    metadata.add_text('prompt',str(prompt))
    for o in history['outputs']:
        for node_id in history['outputs']:
            node_output = history['outputs'][node_id]
            if 'images' in node_output:
                images_output = []
                for image in node_output['images']:
                    image_data = self.get_image(image['filename'], image['subfolder'], image['type'])
                    images_output.append(image_data)
            output_images[node_id] = images_output

    return output_images, metadata

  def draw_image(self, item, draw = True):
    
    filepath = f'{self.savepath}/output_row_{item.row:03d}_col_{item.col:03d}.png'
    print(f'Drawing {filepath} ...')
    
    if draw:
      # positive prompt
      self.json_prompt["6"]["inputs"]["text"] = item.prompt
      # negative prompt
      self.json_prompt["7"]["inputs"]["text"] = item.negative_prompt
      
      self.json_prompt["3"]["inputs"]["seed"] = 333
      self.json_prompt["3"]["inputs"]["cfg"] = 7
      self.json_prompt["3"]["inputs"]["sampler_name"] = "dpm_2"
      self.json_prompt["3"]["inputs"]["scheduler"] = "karras"
      
      self.json_prompt["4"]["inputs"]["ckpt_name"] = "sd_xl_base_1.0.safetensors"
      
      self.json_prompt["5"]["inputs"]["height"] = "1024"
      self.json_prompt["5"]["inputs"]["width"] = "1024"
      
      ws = websocket.WebSocket()
      ws.connect(f"ws://{self.server}/ws?clientId={self.client_id}")
      images, metadata = self.get_images(ws, self.json_prompt)
      ws.close()
      for node_id in images:
        for image_data in images[node_id]:
          image = Image.open(io.BytesIO(image_data))
          image.save(filepath, pnginfo = metadata)
    return filepath


md = MultiDraw()

@dataclass
class ImageData:
  prompt: str
  negative_prompt: str
  row: int
  col: int
  
rows = (
  'teenage girl full face',
  'still life',
  'grumpy old man',
  'picnic in the park',
  'open-air market',
  'children at play',
  'sailors, boat, stormy sea',
  'post-apocalyptic decaying city',
)

# photorealistic and some other arguments must
# precede the prompt, so are marked by <
# others must follow and are given a 'by ' prefix

columns = (
  '<photorealistic',
  '<pencil drawing',
  '<photorealistic art nouveau',
  '<impressionist',
  'Van Gogh',
  'Rembrandt',
  'Vermeer',
  'Hokusai',
  'R. Crumb',
)

html = """
<html>
<body>
<table>
"""

html += '<tr><td></td>\n'
for coln, column in enumerate(columns):
  if column[0] == '<':
    prompt = f'{column[1:]}'
  else:
    prompt = column
  html += f'<td style="width:64px;font-size:80%;text-align:center">{prompt}</td>\n'
html += '</tr>\n'

for rown,row in enumerate(rows):
  if rown == 0:
    negative_prompt = ''
  else:
    negative_prompt = """
  out of frame, lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature,
  """
  print('')
  html += f"""
  <tr>
  <td style="width:64px;font-size:80%;">{row}</td>
  """
  for coln,column in enumerate(columns):
    if column[0] == '<':
      prompt = f'{column[1:]} {row}'
    else:
      prompt = f'{row} by {column}'
    item = ImageData(prompt,negative_prompt,rown,coln)
    # note an available draw argument: True/False
    # this saves time until prompt
    # debugging is complete
    filepath = md.draw_image(item, False)
    html += f"""
    <td>
      <a href="{filepath}">
        <img src="{filepath}" title="Prompt: {prompt} (click for full-size)" width=90/>
      </a>
    </td>
    """    
  html += '</tr>'


html += """
</table>
</body>
</html>
"""

with open('keyindex.html','w') as f:
  f.write(html)