chameleon/viewer/backend/models/chameleon_distributed.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Chameleon License found in the
# LICENSE file in the root directory of this source tree.

import asyncio
import json
import multiprocessing
import os
import random
import sys
import threading
import time
import traceback
from functools import partial
from typing import Any, Generator, TypeVar

import redis
import redis.asyncio as async_redis
import torch
from tokenizers import Tokenizer

from chameleon.inference.image_tokenizer import ImageTokenizer
from chameleon.inference.loader import load_model
from chameleon.inference.vocab import VocabInfo
from chameleon.viewer.backend.data_types import WSMessageType
from chameleon.viewer.backend.models.abstract_model import (
    DEFAULT_IMAGE_CFG_IMAGE,
    DEFAULT_IMAGE_CFG_TEXT,
    DEFAULT_MULTIMODAL_CFG_IMAGE,
    DEFAULT_MULTIMODAL_CFG_TEXT,
    AbstractMultimodalGenerator,
    MixedSequenceType,
    StreamingImage,
)
from chameleon.viewer.backend.models.chameleon_local import (
    ChameleonForwardMixin,
    ChameleonTokenizationMixin,
)
from chameleon.viewer.backend.utils import get_logger

logger = get_logger(__name__)

START = "START"

T = TypeVar("T")


def find_any(queue_by_id: dict[str, list]) -> str | None:
    for candidate_queue_id, candidate_queue in queue_by_id.items():
        if len(candidate_queue) > 0:
            return candidate_queue_id
    return None


class RedisQueue:
    def __init__(self, redis_client: redis.Redis, name: str, interval: float = 0.1):
        self.redis_client = redis_client
        self.name = name
        self.interval = interval
        self.lock = redis.lock.Lock(redis_client, f"lock_for_{name}")

    def reset(self):
        self.redis_client.set(self.name, json.dumps({}))
        try:
            self.lock.release()
        except redis.lock.LockError:
            pass

    def size(self) -> int:
        maybe_queue_by_id = self.redis_client.get(self.name)
        if maybe_queue_by_id is None:
            return 0
        else:
            return len(json.loads(maybe_queue_by_id))

    def clear(self, queue_id: str):
        with self.lock:
            maybe_queue_by_id = self.redis_client.get(self.name)
            if maybe_queue_by_id is None:
                queue_by_id: dict[str, list] = {}
            else:
                queue_by_id: dict[str, list] = json.loads(maybe_queue_by_id)
            queue_by_id[queue_id] = []
            self.redis_client.set(self.name, json.dumps(queue_by_id))

    def put(self, queue_id: str, value: T):
        logger.debug(
            "Thread %s: Starting PUT(%s) for %s",
            threading.get_ident(),
            self.name,
            queue_id,
        )
        with self.lock:
            maybe_queue_by_id = self.redis_client.get(self.name)
            if maybe_queue_by_id is None:
                queue_by_id: dict[str, list[T]] = {}
            else:
                queue_by_id: dict[str, list[T]] = json.loads(maybe_queue_by_id)

            if queue_id not in queue_by_id:
                queue_by_id[queue_id] = []
            queue_by_id[queue_id] = [value] + queue_by_id[queue_id]
            self.redis_client.set(self.name, json.dumps(queue_by_id))

        logger.debug(
            "Thread %s: Finished PUT(%s) for %s",
            threading.get_ident(),
            self.name,
            queue_id,
        )

    def get(self, queue_id: str | None) -> tuple[str, T]:
        """
        Get the next value in the queue.

        if queue_id is None, will get a value from any queue

        if queue_id is not none, will wait to get a value from a specific queue
        """
        logger.debug(
            "Thread %s: Starting GET(%s) for %s",
            threading.get_ident(),
            self.name,
            queue_id,
        )
        while True:
            with self.lock:
                # Initialization hasn't happened, so wait for it to happen
                maybe_queue_by_id = self.redis_client.get(self.name)
                if maybe_queue_by_id is None:
                    continue
                queue_by_id: dict[str, list[T]] = json.loads(maybe_queue_by_id)
                if queue_id is None:
                    queue_id = find_any(queue_by_id)

                # Ensure a queue_id was found or that it already existed
                if queue_id is not None and queue_id in queue_by_id:
                    queue = queue_by_id[queue_id]
                    if len(queue) == 0:
                        continue
                    value = queue.pop(-1)
                    # queue is mutated and queue_by_id references it, so this works
                    self.redis_client.set(self.name, json.dumps(queue_by_id))
                    logger.debug(
                        "Thread %s: Finished GET(%s) for %s",
                        threading.get_ident(),
                        self.name,
                        queue_id,
                    )
                    return queue_id, value
            time.sleep(self.interval)


class AsyncRedisQueue:
    def __init__(
        self, redis_client: async_redis.Redis, name: str, interval: float = 0.1
    ) -> None:
        self.redis_client = redis_client
        self.name = name
        self.interval = interval
        self.lock = async_redis.lock.Lock(redis_client, f"lock_for_{name}")

    async def reset(self):
        await self.redis_client.set(self.name, json.dumps({}))
        try:
            await self.lock.release()
        except async_redis.lock.LockError:
            pass

    async def size(self) -> int:
        maybe_queue_by_id = await self.redis_client.get(self.name)
        if maybe_queue_by_id is None:
            return 0
        else:
            return len(json.loads(maybe_queue_by_id))

    async def clear(self, queue_id: str):
        logger.debug(
            "ASYNC Thread %s: Starting CLEAR(%s) for %s",
            threading.get_ident(),
            self.name,
            queue_id,
        )
        async with self.lock:
            maybe_queue_by_id = await self.redis_client.get(self.name)
            if maybe_queue_by_id is None:
                queue_by_id: dict[str, list] = {}
            else:
                queue_by_id: dict[str, list] = json.loads(maybe_queue_by_id)
            queue_by_id[queue_id] = []
            await self.redis_client.set(self.name, json.dumps(queue_by_id))

        logger.debug(
            "ASYNC Thread %s: Finished CLEAR(%s) for %s",
            threading.get_ident(),
            self.name,
            queue_id,
        )

    async def put(self, queue_id: str, value: T):
        logger.debug(
            "ASYNC Thread %s: Starting PUT(%s) for %s",
            threading.get_ident(),
            self.name,
            queue_id,
        )

        async with self.lock:
            maybe_queue_by_id = await self.redis_client.get(self.name)
            if maybe_queue_by_id is None:
                queue_by_id: dict[str, list[T]] = {}
            else:
                queue_by_id: dict[str, list[T]] = json.loads(maybe_queue_by_id)

            if queue_id not in queue_by_id:
                queue_by_id[queue_id] = []
            queue_by_id[queue_id] = [value] + queue_by_id[queue_id]
            await self.redis_client.set(self.name, json.dumps(queue_by_id))

        logger.debug(
            "ASYNC Thread %s: Finished PUT(%s) for %s",
            threading.get_ident(),
            self.name,
            queue_id,
        )

    async def get(self, queue_id: str | None):
        """
        Get the next value in the queue.

        if queue_id is None, will get a value from any queue

        if queue_id is not none, will wait to get a value from a specific queue
        """
        logger.debug(
            "ASYNC Thread %s: Starting GET(%s) for %s",
            threading.get_ident(),
            self.name,
            queue_id,
        )
        while True:
            async with self.lock:
                maybe_queue_by_id = await self.redis_client.get(self.name)
                if maybe_queue_by_id is None:
                    continue
                queue_by_id: dict[str, list[T]] = json.loads(maybe_queue_by_id)
                if queue_id is None:
                    queue_id = find_any(queue_by_id)

                # Ensure a queue_id was found or that it already existed
                if queue_id is not None and queue_id in queue_by_id:
                    queue: list = queue_by_id[queue_id]
                    if len(queue) == 0:
                        continue
                    value = queue.pop(-1)
                    # queue is mutated and queue_by_id references it, so this works
                    await self.redis_client.set(self.name, json.dumps(queue_by_id))
                    logger.debug(
                        "ASYNC Thread %s: Finished GET(%s) for %s",
                        threading.get_ident(),
                        self.name,
                        queue_id,
                    )
                    return queue_id, value
            await asyncio.sleep(self.interval)


class AsyncRedisCounter:
    def __init__(self, redis_client: async_redis.Redis, name: str) -> None:
        self.redis_client = redis_client
        self.name = name
        self.lock = async_redis.lock.Lock(redis_client, f"lock_for_{name}")

    async def reset(self) -> int:
        try:
            await self.lock.release()
        except async_redis.lock.LockError:
            pass
        await self.redis_client.set(self.name, 0)

    async def add(self, n: int) -> int:
        async with self.lock:
            current_val = await self.redis_client.get(self.name)
            if current_val is None:
                current_val = 0
            else:
                current_val = int(current_val)
            new_val = current_val + n
            await self.redis_client.set(self.name, new_val)
            return new_val

    async def sub(self, n: int) -> int:
        async with self.lock:
            current_val = await self.redis_client.get(self.name)
            if current_val is None:
                raise ValueError("Invalid sub counter when counter does not exist")
            current_val = int(current_val)
            if current_val <= 0:
                raise ValueError("Invalid sub counter to counter that is already zero")
            new_val = current_val - n
            await self.redis_client.set(self.name, new_val)
            return new_val

    async def count(self) -> int:
        value = await self.redis_client.get(self.name)
        if value is None:
            return 0
        else:
            return int(value)


def distributed_workers(
    model_args: dict,
    master_address: str,
    master_port: str,
    world_size: int,
    rank: int,
    redis_port: int,
    worker_queues: dict[int, multiprocessing.Queue],
) -> None:
    redis_client = redis.Redis("redis", redis_port)
    request_queue = RedisQueue(redis_client, "request")
    response_queue = RedisQueue(redis_client, "response")

    os.environ["MASTER_ADDR"] = master_address
    os.environ["MASTER_PORT"] = str(master_port)

    torch.set_default_tensor_type("torch.cuda.FloatTensor")

    torch.distributed.init_process_group("nccl", rank=rank, world_size=world_size)
    assert rank == torch.distributed.get_rank()

    torch.cuda.set_device(rank)

    is_coord = rank == 0

    worker = ChameleonWorker(
        rank=rank,
        model_path=model_args["model_path"],
        tokenizer_path=model_args["tokenizer_path"],
        additional_eos_tokens=model_args["additional_eos_tokens"],
    )
    worker_id = id(worker)
    logger.info("Rank %s, master_port=%s worker=%s", rank, master_port, worker_id)

    step = 0
    while True:
        step += 1
        redis_client.set(f"status_rank_{rank}", "Pre-coordinator sync")
        if is_coord:
            distributed_objs = [request_queue.get(None)]
            logger.info("Objects from queue: %s", distributed_objs)
            for worker_rank in range(1, world_size):
                worker_message = {"message": START, "src": rank, "dst": worker_rank}
                logger.info("Rank %s Sending: %s", rank, worker_message)
                worker_queues[worker_rank].put(worker_message)
        else:
            distributed_objs = [None]
            logger.info("Rank %s worker %s waiting for rank 0", rank, worker_id)
            message_from_rank_0 = worker_queues[rank].get()
            logger.info(
                "Received message from rank 0 in rank %s: %s", rank, message_from_rank_0
            )
            if message_from_rank_0["message"] != START:
                raise ValueError(
                    f"Unexpected message from rank 0: {message_from_rank_0['message']}"
                )
        redis_client.set(f"status_rank_{rank}", "Post-coordinator sync")

        try:
            logger.info(
                "Broadcast Starting: Rank %s, worker %s, step %s",
                rank,
                worker_id,
                step,
            )
            redis_client.set(f"status_rank_{rank}", "Pre-torch sync")
            torch.distributed.broadcast_object_list(distributed_objs, src=0)
            redis_client.set(f"status_rank_{rank}", "Post-torch sync")
            logger.info(
                "Broadcast Complete: Rank %s, worker %s, step %s",
                rank,
                worker_id,
                step,
            )
        except RuntimeError as e:
            logger.error(
                "Rank %s, worker %s, step %s, Error detected in torch broadcast: %s",
                rank,
                worker_id,
                step,
                str(e),
            )
            raise

        logger.info("rank %s, objs %s", rank, distributed_objs)
        queue_id, data = distributed_objs[0]
        mode = data.pop("mode")
        request_id = data.pop("request_id")
        assert queue_id == request_id
        tokenized_prompt = data.pop("tokenized_prompt")
        try:
            match mode:
                case WSMessageType.GENERATE_TEXT:
                    generator_fn = partial(
                        worker._generate_text_streaming, tokenized_prompt, **data
                    )
                case WSMessageType.GENERATE_IMAGE:
                    generator_fn = partial(
                        worker._generate_image_streaming, tokenized_prompt, **data
                    )
                case WSMessageType.GENERATE_MULTIMODAL:
                    generator_fn = partial(
                        worker._generate_multimodal_streaming, tokenized_prompt, **data
                    )
                case _:
                    logger.error(
                        "Encountered unknown mode, crashing the program: %s", mode
                    )
                    response_queue.put(
                        queue_id, {"error": True, "final": True, "message": mode}
                    )
                    raise ValueError("Unknown mode")
            logger.info("Rank: %s, Processing request: %s", rank, request_id)
            i = 0
            redis_client.set(f"status_rank_{rank}", "Pre-generate")
            for output in generator_fn():
                i += 1
                if is_coord:
                    response = {"final": False, "output": output, "error": False}
                    logger.info(
                        "Rank: %s, Adding to response queue: %.100s",
                        rank,
                        response,
                    )
                    redis_client.set(f"status_rank_{rank}", f"Generate Pre Put {i}")
                    response_queue.put(queue_id, response)
                    redis_client.set(f"status_rank_{rank}", f"Generate Post Put {i}")
                else:
                    redis_client.set(f"status_rank_{rank}", f"Generate {i}")
                redis_client.set(f"step_on_rank_{rank}", i)
            redis_client.set(f"status_rank_{rank}", "Post-generate")
            if is_coord:
                logger.info("Rank: %s, Adding final result to output queue", rank)
                response_queue.put(queue_id, {"final": True, "error": False})
        except torch.cuda.OutOfMemoryError as e:
            logger.error("Encountered OOM, crashing the program: %s", e)
            response_queue.put(
                queue_id, {"error": True, "final": True, "message": str(e)}
            )
            crash_program()
        except RuntimeError as e:
            message = str(e)
            if "CUDA" in message:
                logger.error("Encountered CUDA error, crashing the program: %s", e)
                response_queue.put(
                    queue_id, {"error": True, "final": True, "message": str(e)}
                )
                crash_program()
            else:
                logger.error(
                    "Encountered unexpected runtime error, crashing the program: %s %s",
                    e,
                    traceback.format_exc(),
                )
                response_queue.put(
                    queue_id, {"error": True, "final": True, "message": str(e)}
                )
                crash_program()
        except Exception as e:
            logger.error(
                "Encountered unexpected exception: %s %s",
                str(e),
                traceback.format_exc(),
            )
            response_queue.put(
                queue_id, {"error": True, "final": True, "message": str(e)}
            )
            crash_program()


class ChameleonWorker(ChameleonForwardMixin):
    def __init__(
        self,
        *,
        rank: int,
        model_path: str,
        tokenizer_path: str,
        additional_eos_tokens: list[str] | None,
    ) -> None:
        self.rank = rank
        self.model_path = model_path
        self.additional_eos_tokens = additional_eos_tokens
        torch.set_default_device(f"cuda:{rank}")
        self.model = load_model(model_path, rank)
        self.tokenizer = Tokenizer.from_file(str(tokenizer_path))
        self.vocab = VocabInfo(json.load(open(tokenizer_path))["model"]["vocab"])
        logger.info(
            "Rank: %s, Model loaded in worker_obj: %s",
            rank,
            id(self),
        )


def crash_program() -> None:
    logger.error(
        "Crashing the program as instructed, likely due to distributed worker failures"
    )
    sys.exit(1)


class ChameleonDistributedGenerator(AbstractMultimodalGenerator, ChameleonTokenizationMixin):
    def __init__(
        self,
        *,
        world_size: int,
        model_path: str,
        master_port: int,
        tokenizer_path: str,
        vqgan_config_path: str,
        vqgan_ckpt_path: str | None = None,
        master_address: str = "0.0.0.0",
        additional_eos_tokens: list[str] | None = None,
        redis_port: int | None = None,
    ) -> None:
        self.master_port = master_port
        self.master_address = master_address
        self.additional_eos_tokens = additional_eos_tokens
        logger.info("Loading tokenizer...")
        tokenizer_path = tokenizer_path
        self.tokenizer = Tokenizer.from_file(str(tokenizer_path))
        self.vocab = VocabInfo(json.load(open(tokenizer_path))["model"]["vocab"])

        logger.info("Loading VQGAN...")
        self.image_tokenizer = ImageTokenizer(vqgan_config_path, vqgan_ckpt_path)
        self.redis_port = redis_port
        self.redis_pool = async_redis.ConnectionPool.from_url(
            f"redis://redis:{redis_port}"
        )
        self.redis_client = async_redis.Redis.from_pool(self.redis_pool)
        self.request_queue = AsyncRedisQueue(self.redis_client, "request")
        self.response_queue = AsyncRedisQueue(self.redis_client, "response")
        self.worker_queues: dict[int, multiprocessing.Queue] = {
            rank: multiprocessing.Queue() for rank in range(world_size)
        }
        self.procs: list[multiprocessing.Process] = []
        model_args = {
            "model_path": model_path,
            "master_address": master_address,
            "master_port": master_port,
            "tokenizer_path": tokenizer_path,
            "additional_eos_tokens": additional_eos_tokens,
        }
        logger.info("Launching paralle model with world_size=%s", world_size)
        for i in range(world_size):
            proc = multiprocessing.Process(
                target=distributed_workers,
                args=(
                    model_args,
                    master_address,
                    master_port,
                    world_size,
                    i,
                    self.redis_port,
                    self.worker_queues,
                ),
                daemon=True,
            )
            self.procs.append(proc)
            proc.start()

    def check_error(self, output: dict) -> None:
        if output["error"]:
            import sys
            print(f"check_error({output})", file=sys.stderr)
            self.kill_procs()
            logger.error(
                "COORDINATOR: Encountered error in managed processes, exiting: %s",
                output,
            )
            crash_program()

    def __del__(self) -> None:
        self.kill_procs(error=False)

    def kill_procs(self, error: bool = True) -> None:
        if error:
            log_fn = logger.error
        else:
            log_fn = logger.info
        log_fn("Error encountered, killing worker procs: %s", self.procs)
        for p in self.procs:
            try:
                log_fn("Killing: %s", p)
                p.kill()
            except:
                log_fn("Encountered issue killing process and ignoring: %s", p)

    # ALLOW_ANY(get_next_output.return)
    async def get_next_output(self, request_id: str) -> Any:
        logger.info("Waiting for response for request_id=%s", request_id)
        queue_id, output = await self.response_queue.get(request_id)
        assert queue_id == request_id
        return output

    async def generate_text_streaming(
        self,
        prompt: MixedSequenceType,
        max_gen_tokens: int = 256,
        temp: float = 1.0,
        top_p: float = 0.8,
        repetition_penalty: float = 1.2,
        seed: int | None = None,
        debug: dict | None = None,
    ) -> Generator[str, None, None]:
        tokenized_prompt = self.tokens_from_inputs(prompt)
        request_id = f"request_{random.randint(100_000, 200_000)}"
        if seed is None:
            seed = random.randint(1, 2048)
            if debug is not None:
                debug["seed"] = seed
        if len(tokenized_prompt) > (4096 - 3):
            yield "ERROR: Your input exceeds the model's context length of 4096. Note that images consume 1024 tokens whether in input or output."
            return
        assert not isinstance(tokenized_prompt, torch.Tensor)
        request = {
            "mode": WSMessageType.GENERATE_TEXT.value,
            "request_id": request_id,
            "tokenized_prompt": tokenized_prompt,
            "max_gen_tokens": max_gen_tokens,
            "temp": temp,
            "top_p": top_p,
            "repetition_penalty": repetition_penalty,
            "seed": seed,
        }
        logger.info(
            "Sending request_id=%s: %s",
            request_id,
            request,
        )
        await asyncio.gather(
            self.request_queue.clear(request_id),
            self.response_queue.clear(request_id),
        )
        logger.info("Cleared request/response queue for %s", request_id)
        await self.request_queue.put(request_id, request)
        logger.info("Sent request to coordinator %s", request_id)
        try:
            while True:
                output = await self.get_next_output(request_id)
                logger.info("Received response for %s", request_id)
                self.check_error(output)
                if output["final"]:
                    break

                n_outs = len(output["output"])
                if n_outs != 1:
                    logger.error(
                        "Encountered unexpected number of %s arguments in: %s",
                        n_outs,
                        output["output"],
                    )
                tokens = output["output"]
                assert not isinstance(tokens, torch.Tensor)
                logger.info("output info: type=%s, value=%.20s", type(tokens), tokens)
                yield self.tokenizer.decode(tokens)
        finally:
            logger.info("Cleaning up queues in request_id=%s", request_id)
            await asyncio.gather(
                self.request_queue.clear(request_id),
                self.response_queue.clear(request_id),
            )
            logger.info("Completed cleaning for request_id=%s", request_id)

    async def generate_image_streaming(
        self,
        prompt: MixedSequenceType,
        temp: float = 1.0,
        top_p: float = 0.8,
        cfg_image_weight: float = DEFAULT_IMAGE_CFG_IMAGE,
        cfg_text_weight: float = DEFAULT_IMAGE_CFG_TEXT,
        yield_every_n: int = 32,
        debug: dict | None = None,
        seed: int | None = None,
    ) -> Generator[StreamingImage, None, None]:
        tokenized_prompt = self.tokens_from_inputs(prompt)
        tokenized_prompt.append(self.vocab.begin_image)
        assert not isinstance(tokenized_prompt, torch.Tensor)
        request_id = f"request_{random.randint(100_000, 200_000)}"
        if seed is None:
            seed = random.randint(1, 2048)
            if debug is not None:
                debug["seed"] = seed
        if len(tokenized_prompt) > (4096 - 3 - 1024):
            yield "ERROR: Your input exceeds the model's context length of 4096. Note that images consume 1024 tokens whether in input or output."
            return
        request = {
            "mode": WSMessageType.GENERATE_IMAGE.value,
            "request_id": request_id,
            "tokenized_prompt": tokenized_prompt,
            "cfg_image_weight": cfg_image_weight,
            "cfg_text_weight": cfg_text_weight,
            "yield_every_n": yield_every_n,
            "temp": temp,
            "top_p": top_p,
            "seed": seed,
        }
        logger.info(
            "Sending request_id=%s: %s",
            request_id,
            request,
        )
        await asyncio.gather(
            self.request_queue.clear(request_id),
            self.response_queue.clear(request_id),
        )
        logger.info("Cleared request/response queue for %s", request_id)
        await self.request_queue.put(request_id, request)
        logger.info("Sent request to coordinator %s", request_id)
        try:
            while True:
                output = await self.get_next_output(request_id)
                logger.info("Received response for %s", request_id)
                self.check_error(output)
                if output["final"]:
                    break
                n_outs = len(output["output"])
                if n_outs != 2:
                    logger.error(
                        "Encountered unexpected number of %s arguments in: %s",
                        n_outs,
                        output["output"],
                    )
                tokens, final = output["output"]
                assert not isinstance(tokens, torch.Tensor)
                yield StreamingImage(
                    image=self.pillow_from_bpe_tokens(torch.tensor(tokens)), final=final
                )
        finally:
            logger.info("Cleaning up queues in request_id=%s", request_id)
            await asyncio.gather(
                self.request_queue.clear(request_id),
                self.response_queue.clear(request_id),
            )
            logger.info("Completed cleaning for request_id=%s", request_id)

    async def generate_multimodal_streaming(
        self,
        prompt: MixedSequenceType,
        temp: float = 1.0,
        top_p: float = 0.8,
        cfg_image_weight: float = DEFAULT_MULTIMODAL_CFG_IMAGE,
        cfg_text_weight: float = DEFAULT_MULTIMODAL_CFG_TEXT,
        yield_every_n: int = 32,
        max_gen_tokens: int = 4096,
        repetition_penalty: float = 1.2,
        suffix_tokens: list[str] | None = None,
        seed: int | None = None,
        debug: dict | None = None,
    ) -> Generator[MixedSequenceType, None, None]:
        tokenized_prompt = self.tokens_from_inputs(prompt, suffix_tokens=suffix_tokens)
        assert not isinstance(tokenized_prompt, torch.Tensor)
        request_id = f"request_{random.randint(100_000, 200_000)}"
        if seed is None:
            seed = random.randint(1, 2048)
            if debug is not None:
                debug["seed"] = seed
        if len(tokenized_prompt) > (4096 - 3):
            yield "ERROR: Your input exceeds the model's context length of 4096. Note that images consume 1024 tokens."
            return

        request = {
            "mode": WSMessageType.GENERATE_MULTIMODAL.value,
            "request_id": request_id,
            "tokenized_prompt": tokenized_prompt,
            "cfg_image_weight": cfg_image_weight,
            "cfg_text_weight": cfg_text_weight,
            "repetition_penalty": repetition_penalty,
            "yield_every_n": yield_every_n,
            "max_gen_tokens": max_gen_tokens,
            "temp": temp,
            "top_p": top_p,
            "seed": seed,
        }
        logger.info(
            "Sending request_id=%s: %s",
            request_id,
            request,
        )
        await asyncio.gather(
            self.request_queue.clear(request_id),
            self.response_queue.clear(request_id),
        )
        logger.info("Cleared request/response queue for %s", request_id)
        await self.request_queue.put(request_id, request)
        logger.info("Sent request to coordinator %s", request_id)
        try:
            while True:
                output = await self.get_next_output(request_id)
                logger.info("Received response for %s", request_id)
                self.check_error(output)
                if output["final"]:
                    break
                n_outs = len(output["output"])
                if n_outs != 3:
                    logger.error(
                        "Encountered unexpected number of %s arguments in: %s",
                        n_outs,
                        output["output"],
                    )
                token_type, tokens, image_is_final = output["output"]
                assert not isinstance(tokens, torch.Tensor)
                match token_type:
                    case "TEXT":
                        yield self.tokenizer.decode(tokens)
                    case "IMAGE":
                        yield StreamingImage(
                            image=self.pillow_from_bpe_tokens(torch.tensor(tokens)),
                            final=image_is_final,
                        )
                    case _:
                        raise ValueError("Unknown token type")
        finally:
            logger.info("Cleaning up queues in request_id=%s", request_id)
            await self.request_queue.clear(request_id)
            await self.response_queue.clear(request_id)