Created
July 1, 2024 10:50
-
-
Save jpfraneto/aca5c52dd3f66afa553c46ff8388caad to your computer and use it in GitHub Desktop.
from lines 287 to 321, you can see the addition of the /api/interpret_post route to the handler of the backend of chameleon. this allows us to talk to this model over an api request, and have it interpret images with text.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# 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 io | |
import socket | |
import subprocess | |
import time | |
from functools import partial | |
import fastapi | |
import PIL | |
import pydantic | |
import redis.asyncio as async_redis | |
import uvicorn | |
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, WebSocketException, HTTPException, UploadFile, File, Form | |
from typing import Optional | |
from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK | |
from chameleon.viewer.backend.data_types import ( | |
Content, | |
ContentType, | |
NoOptionsForComplete, | |
NoOptionsForFull, | |
NoOptionsForPartial, | |
NoOptionsForQueueStatus, | |
WSMessageType, | |
WSMultimodalMessage, | |
) | |
from chameleon.viewer.backend.models.abstract_model import ( | |
AbstractMultimodalGenerator, | |
StreamingImage, | |
) | |
from chameleon.viewer.backend.models.chameleon_distributed import AsyncRedisCounter | |
from chameleon.viewer.backend.utils import get_logger | |
logger = get_logger(__name__) | |
def nvidia_smi() -> str: | |
return subprocess.check_output(["nvidia-smi"], text=True) | |
async def await_generate_message(websocket: WebSocket) -> WSMultimodalMessage: | |
while True: | |
rec_message = await websocket.receive_json() | |
try: | |
maybe_message = WSMultimodalMessage.parse_obj(rec_message) | |
except pydantic.ValidationError: | |
maybe_message = None | |
logger.info("Got invalid message", maybe_message) | |
if maybe_message is not None: | |
return maybe_message | |
async def async_acquire_lock( | |
*, | |
websocket: WebSocket, | |
counter: AsyncRedisCounter, | |
lock: async_redis.lock.Lock, | |
interval=0.1, | |
status_interval=1, | |
hostname: str | None = None, | |
): | |
start = time.time() | |
await counter.add(1) | |
while True: | |
acquired = await lock.acquire(blocking_timeout=interval) | |
if acquired: | |
break | |
elapsed = time.time() - start | |
if elapsed > status_interval: | |
n_requests = await counter.count() | |
message = WSMultimodalMessage( | |
message_type=WSMessageType.QUEUE_STATUS, | |
content=[ | |
Content( | |
content_type=ContentType.TEXT, | |
content=f"n_requests={n_requests}", | |
) | |
], | |
options=NoOptionsForQueueStatus(), | |
debug_info={"hostname": hostname}, | |
).dict() | |
await websocket.send_json(message) | |
start = time.time() | |
await counter.sub(1) | |
COORDINATOR = "coordinator" | |
def web_app( | |
generator: AbstractMultimodalGenerator, | |
debug: bool = True, | |
redis_port: int | None = None, | |
) -> FastAPI: | |
app = FastAPI(debug=debug) | |
if redis_port is None: | |
redis_client = None | |
redis_lock = None | |
queue_counter = None | |
else: | |
redis_client = async_redis.Redis.from_url(f"redis://redis:{redis_port}") | |
redis_lock = async_redis.lock.Lock(redis_client, COORDINATOR) | |
queue_counter = AsyncRedisCounter(redis_client, "count_pending") | |
hostname = socket.gethostname() | |
@app.get("/api/2.0/status") | |
def alive() -> dict: | |
return { | |
"status": "alive", | |
"hostname": hostname, | |
"nvidia-smi": nvidia_smi(), | |
} | |
@app.websocket("/ws/chameleon/v2/{client_id}") | |
async def websocket_chameleon_v2(*, websocket: WebSocket, client_id: str): | |
logger.info("Requested client_id: %s", client_id) | |
await websocket.accept() | |
logger.info("Client opened %s with generator id %s", client_id, id(generator)) | |
try: | |
while True: | |
generate_message = await await_generate_message(websocket) | |
logger.info("Got generate message: %s", str(generate_message)[:300]) | |
parsed_prompt = [] | |
for c in generate_message.content: | |
match c.content_type: | |
case ContentType.TEXT: | |
parsed_prompt.append(c.content) | |
case ContentType.IMAGE: | |
image_parts = c.content.split(",", 1) | |
if len(image_parts) < 2: | |
logger.error( | |
"Encountered invalid image: %s", image_parts | |
) | |
raise WebSocketException( | |
code=fastapi.status.WS_1008_POLICY_VIOLATION, | |
reason=f"Invalid image: {image_parts}", | |
) | |
image_data = image_parts[1] | |
base64_image = base64.b64decode(image_data) | |
image_file = io.BytesIO(base64_image) | |
parsed_prompt.append(PIL.Image.open(image_file)) | |
case _: | |
raise ValueError("Unknown content type") | |
logger.info("Prompt: %s", parsed_prompt) | |
partial_outputs = [] | |
final_contents: list[Content] = [] | |
match generate_message.message_type: | |
case WSMessageType.GENERATE_TEXT: | |
output_generator = generator.generate_text_streaming | |
case WSMessageType.GENERATE_IMAGE: | |
output_generator = generator.generate_image_streaming | |
case WSMessageType.GENERATE_MULTIMODAL: | |
output_generator = generator.generate_multimodal_streaming | |
case _: | |
raise WebSocketException( | |
code=fastapi.status.WS_1008_POLICY_VIOLATION, | |
reason="Unknown message type", | |
) | |
logger.info( | |
"Acquiring lock for client %s generation with options: %s", | |
client_id, | |
generate_message.options, | |
) | |
option_args = generate_message.options.dict() | |
debug_info = {"hostname": hostname} | |
del option_args["message_type"] | |
output_generator = partial( | |
output_generator, | |
**option_args, | |
debug=debug_info, | |
) | |
if redis_lock is not None: | |
await async_acquire_lock( | |
websocket=websocket, | |
lock=redis_lock, | |
hostname=hostname, | |
counter=queue_counter, | |
) | |
await redis_client.set("has_lock", client_id) | |
logger.info( | |
"Starting locked generation for client %s with options: %s", | |
client_id, | |
generate_message.options, | |
) | |
try: | |
async for output_token in output_generator(parsed_prompt): | |
if isinstance(output_token, str): | |
content_type = ContentType.TEXT | |
content = output_token | |
message_type = WSMessageType.PARTIAL_OUTPUT | |
options = NoOptionsForPartial() | |
partial_outputs.extend(output_token) | |
elif isinstance(output_token, StreamingImage): | |
content_type = ContentType.IMAGE | |
image = output_token.image | |
img_io = io.BytesIO() | |
image.save(img_io, format="png") | |
content = ( | |
"data:image/png;base64," | |
+ base64.b64encode(img_io.getvalue()).decode() | |
) | |
if output_token.final: | |
message_type = WSMessageType.FULL_OUTPUT | |
options = NoOptionsForFull() | |
else: | |
message_type = WSMessageType.PARTIAL_OUTPUT | |
options = NoOptionsForPartial() | |
if output_token.final: | |
partial_outputs.append(output_token.image) | |
else: | |
raise ValueError(f"Invalid output_token: {output_token}") | |
message_content = Content( | |
content_type=content_type, content=content | |
) | |
match content_type: | |
case ContentType.TEXT: | |
final_contents.append(message_content) | |
case ContentType.IMAGE: | |
if message_type == WSMessageType.FULL_OUTPUT: | |
final_contents.append(message_content) | |
case _: | |
pass | |
message = WSMultimodalMessage( | |
message_type=message_type, | |
content=[message_content], | |
options=options, | |
debug_info=debug_info, | |
).dict() | |
await websocket.send_json(message) | |
finally: | |
if redis_lock is not None: | |
logger.info( | |
"Attempting release of lock for client %s generation with options: %s", | |
client_id, | |
generate_message.options, | |
) | |
owned = await redis_lock.owned() | |
if owned: | |
await redis_client.set("has_lock", "") | |
try: | |
await redis_lock.release() | |
except async_redis.lock.LockError: | |
pass | |
logger.info( | |
"Released lock for client %s generation with options: %s", | |
client_id, | |
generate_message.options, | |
) | |
await websocket.send_json( | |
WSMultimodalMessage( | |
message_type=WSMessageType.COMPLETE, | |
content=final_contents, | |
options=NoOptionsForComplete(), | |
debug_info=debug_info, | |
).dict() | |
) | |
except WebSocketDisconnect: | |
logger.info("Client disconnected %s", client_id) | |
except ConnectionClosedError: | |
logger.info("Client forced a close %s", client_id) | |
except ConnectionClosedOK: | |
logger.info("Connection closed ok %s", client_id) | |
finally: | |
if redis_lock is not None: | |
logger.info("Checking for client holding lock: %s", client_id) | |
owned = await redis_lock.owned() | |
if owned: | |
try: | |
logger.info("Attempted to release owned lock: %s", client_id) | |
await redis_lock.release() | |
except async_redis.lock.LockError: | |
pass | |
await redis_client.set("has_lock", "") | |
@app.post("/api/interpret_post") | |
async def interpret_post( | |
image: Optional[UploadFile] = File(None), | |
text: str = Form(...), | |
): | |
try: | |
parsed_prompt = [text] | |
logger.info(f"Received text: {text}") | |
if image: | |
image_content = await image.read() | |
image_file = io.BytesIO(image_content) | |
parsed_prompt.append(PIL.Image.open(image_file)) | |
logger.info("Image received and processed") | |
else: | |
logger.info("No image received") | |
logger.info(f"Parsed prompt: {parsed_prompt}") | |
full_interpretation = "" | |
logger.info("Starting generation...") | |
# Use generate_multimodal_streaming but collect all output | |
async for output_token in generator.generate_multimodal_streaming(parsed_prompt): | |
if isinstance(output_token, str): | |
full_interpretation += output_token | |
# We're still ignoring StreamingImage outputs for this use case | |
logger.info(f"Generation complete. Full interpretation length: {len(full_interpretation)}") | |
return {"interpretation": full_interpretation} | |
except Exception as e: | |
logger.error(f"Error in interpret_post: {str(e)}") | |
logger.error(f"Error traceback: {traceback.format_exc()}") | |
raise HTTPException(status_code=500, detail=str(e)) | |
return app | |
def serve( | |
model: AbstractMultimodalGenerator, | |
host: str, | |
port: int, | |
debug: bool = True, | |
redis_port: int | None = None, | |
) -> None: | |
app = web_app(model, debug=debug, redis_port=redis_port) | |
# TODO: convert this to a subprocess call so enable more | |
# uvicorn features like multiple workers | |
uvicorn.run(app, host=host, port=port) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment