303 lines
8.7 KiB
Python
303 lines
8.7 KiB
Python
# This is an example that uses the websockets api to know when a prompt execution is done
|
|
# Once the prompt execution is done it downloads the images using the /history endpoint
|
|
|
|
import io
|
|
import json
|
|
import urllib.parse
|
|
import urllib.request
|
|
import uuid
|
|
from logging import getLogger
|
|
from os import environ, path
|
|
from queue import Queue
|
|
from random import choice, randint
|
|
from threading import Thread
|
|
from typing import List
|
|
|
|
import websocket # NOTE: websocket-client (https://github.com/websocket-client/websocket-client)
|
|
from PIL import Image
|
|
|
|
from adventure.context import broadcast
|
|
from adventure.models.event import (
|
|
ActionEvent,
|
|
GameEvent,
|
|
RenderEvent,
|
|
ReplyEvent,
|
|
ResultEvent,
|
|
StatusEvent,
|
|
)
|
|
|
|
logger = getLogger(__name__)
|
|
|
|
server_address = environ["COMFY_API"]
|
|
client_id = str(uuid.uuid4())
|
|
|
|
|
|
def generate_cfg():
|
|
return randint(5, 8)
|
|
|
|
|
|
def generate_steps():
|
|
return 30
|
|
|
|
|
|
def generate_batches(
|
|
count: int,
|
|
batch_size: int = 3,
|
|
) -> List[int]:
|
|
"""
|
|
Generate count images in batches of at most batch_size.
|
|
"""
|
|
|
|
batches = []
|
|
for i in range(0, count, batch_size):
|
|
batches.append(min(count - i, batch_size))
|
|
|
|
return batches
|
|
|
|
|
|
def queue_prompt(prompt):
|
|
p = {"prompt": prompt, "client_id": client_id}
|
|
data = json.dumps(p).encode("utf-8")
|
|
req = urllib.request.Request("http://{}/prompt".format(server_address), data=data)
|
|
return json.loads(urllib.request.urlopen(req).read())
|
|
|
|
|
|
def get_image(filename, subfolder, folder_type):
|
|
data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
|
|
url_values = urllib.parse.urlencode(data)
|
|
with urllib.request.urlopen(
|
|
"http://{}/view?{}".format(server_address, url_values)
|
|
) as response:
|
|
return response.read()
|
|
|
|
|
|
def get_history(prompt_id):
|
|
with urllib.request.urlopen(
|
|
"http://{}/history/{}".format(server_address, prompt_id)
|
|
) as response:
|
|
return json.loads(response.read())
|
|
|
|
|
|
def get_images(ws, prompt):
|
|
prompt_id = 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 # Execution is done
|
|
else:
|
|
continue # previews are binary data
|
|
|
|
history = get_history(prompt_id)[prompt_id]
|
|
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 = get_image(
|
|
image["filename"], image["subfolder"], image["type"]
|
|
)
|
|
images_output.append(image_data)
|
|
output_images[node_id] = images_output
|
|
|
|
return output_images
|
|
|
|
|
|
def generate_image_tool(prompt, count, size="landscape"):
|
|
output_paths = []
|
|
for i, count in enumerate(generate_batches(count)):
|
|
results = generate_images(prompt, count, size, prefix=f"output-{i}")
|
|
output_paths.extend(results)
|
|
|
|
return output_paths
|
|
|
|
|
|
sizes = {
|
|
"landscape": (1024, 768),
|
|
"portrait": (768, 1024),
|
|
"square": (768, 768),
|
|
}
|
|
|
|
|
|
def generate_images(
|
|
prompt: str, count: int, size="landscape", prefix="output"
|
|
) -> List[str]:
|
|
cfg = generate_cfg()
|
|
width, height = sizes.get(size, (512, 512))
|
|
steps = generate_steps()
|
|
seed = randint(0, 10000000)
|
|
checkpoint = choice(["diffusion-sdxl-dynavision-0-5-5-7.safetensors"])
|
|
logger.info(
|
|
"generating %s images at %s by %s with prompt: %s", count, width, height, prompt
|
|
)
|
|
|
|
# parsing here helps ensure the template emits valid JSON
|
|
prompt_workflow = {
|
|
"3": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"cfg": cfg,
|
|
"denoise": 1,
|
|
"latent_image": ["5", 0],
|
|
"model": ["4", 0],
|
|
"negative": ["7", 0],
|
|
"positive": ["6", 0],
|
|
"sampler_name": "euler_ancestral",
|
|
"scheduler": "normal",
|
|
"seed": seed,
|
|
"steps": steps,
|
|
},
|
|
},
|
|
"4": {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": checkpoint},
|
|
},
|
|
"5": {
|
|
"class_type": "EmptyLatentImage",
|
|
"inputs": {"batch_size": count, "height": height, "width": width},
|
|
},
|
|
"6": {
|
|
"class_type": "smZ CLIPTextEncode",
|
|
"inputs": {
|
|
"text": prompt,
|
|
"parser": "compel",
|
|
"mean_normalization": True,
|
|
"multi_conditioning": True,
|
|
"use_old_emphasis_implementation": False,
|
|
"with_SDXL": False,
|
|
"ascore": 6,
|
|
"width": width,
|
|
"height": height,
|
|
"crop_w": 0,
|
|
"crop_h": 0,
|
|
"target_width": width,
|
|
"target_height": height,
|
|
"text_g": "",
|
|
"text_l": "",
|
|
"smZ_steps": 1,
|
|
"clip": ["4", 1],
|
|
},
|
|
},
|
|
"7": {"class_type": "CLIPTextEncode", "inputs": {"text": "", "clip": ["4", 1]}},
|
|
"8": {
|
|
"class_type": "VAEDecode",
|
|
"inputs": {"samples": ["3", 0], "vae": ["4", 2]},
|
|
},
|
|
"9": {
|
|
"class_type": "SaveImage",
|
|
"inputs": {"filename_prefix": prefix, "images": ["8", 0]},
|
|
},
|
|
}
|
|
|
|
logger.debug("Connecting to Comfy API at %s", server_address)
|
|
ws = websocket.WebSocket()
|
|
ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id))
|
|
images = get_images(ws, prompt_workflow)
|
|
|
|
results = []
|
|
for node_id in images:
|
|
for image_data in images[node_id]:
|
|
image = Image.open(io.BytesIO(image_data))
|
|
results.append(image)
|
|
|
|
paths: List[str] = []
|
|
for j, image in enumerate(results):
|
|
# TODO: replace with environment variable
|
|
image_path = path.join("/home/ssube/adventure-images", f"{prefix}-{j}.png")
|
|
with open(image_path, "wb") as f:
|
|
image_bytes = io.BytesIO()
|
|
image.save(image_bytes, format="PNG")
|
|
f.write(image_bytes.getvalue())
|
|
|
|
paths.append(image_path)
|
|
|
|
return paths
|
|
|
|
|
|
def prompt_from_event(event: GameEvent) -> str | None:
|
|
if isinstance(event, ActionEvent):
|
|
if event.item:
|
|
return f"{event.actor.name} uses the {event.item.name}. {event.item.description}. {event.actor.description}. {event.room.description}."
|
|
|
|
return f"{event.actor.name} {event.action}. {event.actor.description}. {event.room.description}."
|
|
|
|
if isinstance(event, ReplyEvent):
|
|
return event.text
|
|
|
|
if isinstance(event, ResultEvent):
|
|
return f"{event.result}. {event.actor.description}. {event.room.description}."
|
|
|
|
if isinstance(event, StatusEvent):
|
|
if event.room:
|
|
if event.actor:
|
|
return f"{event.text}. {event.actor.description}. {event.room.description}."
|
|
|
|
return f"{event.text}. {event.room.description}."
|
|
|
|
return event.text
|
|
|
|
return None
|
|
|
|
|
|
def prefix_from_event(event: GameEvent) -> str:
|
|
if isinstance(event, ActionEvent):
|
|
return (
|
|
f"{event.actor.name}-{event.action}-{event.item.name if event.item else ''}"
|
|
)
|
|
|
|
if isinstance(event, ReplyEvent):
|
|
return f"{event.actor.name}-reply"
|
|
|
|
if isinstance(event, ResultEvent):
|
|
return f"{event.actor.name}-result"
|
|
|
|
if isinstance(event, StatusEvent):
|
|
return "status"
|
|
|
|
return "unknown"
|
|
|
|
|
|
# requests to generate images for game events
|
|
render_queue: Queue[GameEvent] = Queue()
|
|
|
|
|
|
def render_loop():
|
|
while True:
|
|
event = render_queue.get()
|
|
prompt = prompt_from_event(event)
|
|
if prompt:
|
|
logger.info("rendering prompt for event %s: %s", event, prompt)
|
|
prefix = prefix_from_event(event)
|
|
image_paths = generate_images(prompt, 2, prefix=prefix)
|
|
broadcast(RenderEvent(paths=image_paths, source=event))
|
|
else:
|
|
logger.warning("no prompt for event %s", event)
|
|
|
|
|
|
def render_event(event: GameEvent):
|
|
render_queue.put(event)
|
|
|
|
|
|
render_thread = None
|
|
|
|
|
|
def launch_render():
|
|
global render_thread
|
|
|
|
render_thread = Thread(target=render_loop, daemon=True)
|
|
render_thread.start()
|
|
|
|
return [render_thread]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
paths = generate_images(
|
|
"A painting of a beautiful sunset over a calm lake", 3, "landscape"
|
|
)
|
|
logger.info("Generated %d images: %s", len(paths), paths)
|