# 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 base64
import os
import threading
import time
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
import click
import torch
from flask import Flask, request
from flask_socketio import SocketIO
from chameleon.inference.chameleon import ChameleonInferenceModel, Options, TokenManager
@dataclass
class Request:
room: str
key: str
options: dict[str, int | float | bool]
prompt_ui: list[dict]
def convert_options(ui_options: dict) -> Options:
txt = None
if ui_options["enable-text"]:
txt = Options.Text(
repetition_penalty=ui_options["text-rep-penalty"],
temp=ui_options["text-temp"],
top_p=ui_options["text-top-p"],
)
img = None
if ui_options["enable-image"]:
img = Options.Image(
cfg=Options.Image.CFG(
guidance_scale_image=ui_options["img-cfg-gsimage"],
guidance_scale_text=ui_options["img-cfg-gstext"],
),
temp=ui_options["img-temp"],
top_p=ui_options["img-top-p"],
)
return Options(
max_seq_len=ui_options["max-seq-len"],
max_gen_len=ui_options["max-gen-len"],
seed=ui_options["seed"],
txt=txt,
img=img,
)
class UIDecoder:
class State(Enum):
TXT = 1
IMG = 2
IMG_END = 3
def __init__(self, token_manager: TokenManager):
self.token_manager = token_manager
self.state = UIDecoder.State.TXT
self.image_builder = []
self.image_yield_every_n = 32
self.image_has_updated = False
def _image_progress(self) -> dict:
self.image_has_updated = False
png = self.token_manager.png_from_bpe_tokens(torch.cat(self.image_builder))
return {
"type": "image",
"value": "data:image/png;base64," + base64.b64encode(png).decode(),
}
def next(self, gpu_token: torch.LongTensor) -> dict | None:
if self.state == UIDecoder.State.TXT:
cpu_tok = gpu_token.item()
if cpu_tok == self.token_manager.vocab.begin_image:
self.state = UIDecoder.State.IMG
return {"type": "image_start"}
return {
"type": "text",
"value": self.token_manager.tokenizer.decode([cpu_tok]),
}
elif self.state == UIDecoder.State.IMG:
self.image_builder.append(gpu_token)
self.image_has_updated = True
if len(self.image_builder) == 1024:
self.state = UIDecoder.State.IMG_END
if len(self.image_builder) % self.image_yield_every_n == 0:
return self._image_progress()
elif self.state == UIDecoder.State.IMG_END:
# assert gpu_token == end_image
self.state = UIDecoder.State.TXT
progress = self._image_progress() if self.image_has_updated else None
self.image_builder = []
return progress
@dataclass
class State:
room_keys: dict[str, set[str]]
pending_requests: list[Request]
cond: threading.Condition
def __enter__(self, *args, **kwargs):
self.cond.__enter__(*args, **kwargs)
return self
def __exit__(self, *args, **kwargs):
self.cond.__exit__(*args, **kwargs)
return self
GlobalState = State(room_keys={}, pending_requests=[], cond=threading.Condition())
app = Flask(__name__)
socketio = SocketIO(app, max_http_buffer_size=16 * 1024 * 1024)
@app.route("/")
def index():
with open(Path(__file__).parent / "miniviewer.html") as f:
return f.read()
@socketio.on("disconnect")
def handle_disconnect():
with GlobalState as state:
try:
del state.room_keys[request.sid]
except KeyError:
pass
@socketio.on("cancel")
def handle_cancel(key):
with GlobalState as state:
try:
state.room_keys[request.sid].remove(key)
except KeyError:
pass
@socketio.on("generate")
def handle_generate(key, options, prompt_ui):
with GlobalState as state:
if request.sid not in state.room_keys:
state.room_keys[request.sid] = set()
state.room_keys[request.sid].add(key)
state.pending_requests.append(Request(request.sid, key, options, prompt_ui))
state.cond.notify_all()
def generation_thread(model: ChameleonInferenceModel):
while True:
with GlobalState as state:
state.cond.wait_for(lambda: state.pending_requests)
req = state.pending_requests.pop(0)
start = time.time()
ui_decoder = UIDecoder(model.token_manager)
options = convert_options(req.options)
if not options.txt:
progress = ui_decoder.next(
torch.tensor([model.token_manager.vocab.begin_image])
)
socketio.emit(
"progress",
{"key": req.key, **progress},
room=req.room,
)
for token in model.stream(
prompt_ui=req.prompt_ui,
options=options,
):
with GlobalState as state:
if req.key not in state.room_keys.get(req.room, {}):
break
if progress := ui_decoder.next(token.id):
socketio.emit(
"progress",
{"key": req.key, **progress},
room=req.room,
)
timing = time.time() - start
socketio.emit(
"progress",
{"key": req.key, "type": "done", "value": timing},
room=req.room,
)
def queue_position_thread():
local_pending_requests = []
while True:
with GlobalState as state:
state.cond.wait_for(
lambda: local_pending_requests != state.pending_requests
)
local_pending_requests = state.pending_requests[:]
for i, req in enumerate(local_pending_requests):
progress = {
"type": "queue",
"key": req.key,
"value": i + 1,
}
socketio.emit("progress", progress, room=req.room)
@click.command()
@click.option("--data-path", type=click.Path(), default="./data")
@click.option(
"--model-size", type=click.Choice(["7b", "30b"], case_sensitive=False), default="7b"
)
def main(data_path, model_size):
data_path = Path(data_path)
model_path = str(data_path / "models" / model_size)
tokenizer_path = str(data_path / "tokenizer/text_tokenizer.json")
vqgan_cfg_path = str(data_path / "tokenizer/vqgan.yaml")
vqgan_ckpt_path = str(data_path / "tokenizer/vqgan.ckpt")
if not os.path.exists(model_path):
raise ValueError(
"Model not found. Did you run python -m chameleon.download_data {PRESIGNED_URL}"
)
cm3v2_inference_model = ChameleonInferenceModel(
model_path, tokenizer_path, vqgan_cfg_path, vqgan_ckpt_path
)
threading.Thread(
target=generation_thread,
args=(cm3v2_inference_model,),
daemon=True,
).start()
threading.Thread(target=queue_position_thread, daemon=True).start()
socketio.run(app, debug=False)
if __name__ == "__main__":
main()