diff --git a/demo_w_camera/config.py b/demo_w_camera/config.py new file mode 100644 index 0000000..4121d5b --- /dev/null +++ b/demo_w_camera/config.py @@ -0,0 +1,109 @@ +from typing import NamedTuple +import argparse +import os + + +class Args(NamedTuple): + host: str + port: int + reload: bool + mode: str + max_queue_size: int + timeout: float + safety_checker: bool + taesd: bool + ssl_certfile: str + ssl_keyfile: str + debug: bool + acceleration: str + engine_dir: str + + def pretty_print(self): + print("\n") + for field, value in self._asdict().items(): + print(f"{field}: {value}") + print("\n") + + +MAX_QUEUE_SIZE = int(os.environ.get("MAX_QUEUE_SIZE", 0)) +TIMEOUT = float(os.environ.get("TIMEOUT", 0)) +SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", None) == "True" +USE_TAESD = os.environ.get("USE_TAESD", "True") == "True" +ENGINE_DIR = os.environ.get("ENGINE_DIR", "engines") +ACCELERATION = os.environ.get("ACCELERATION", "xformers") + +default_host = os.getenv("HOST", "0.0.0.0") +default_port = int(os.getenv("PORT", "7860")) +default_mode = os.getenv("MODE", "default") + +parser = argparse.ArgumentParser(description="Run the app") +parser.add_argument("--host", type=str, default=default_host, help="Host address") +parser.add_argument("--port", type=int, default=default_port, help="Port number") +parser.add_argument("--reload", action="store_true", help="Reload code on change") +parser.add_argument( + "--mode", type=str, default=default_mode, help="App Inferece Mode: txt2img, img2img" +) +parser.add_argument( + "--max-queue-size", + dest="max_queue_size", + type=int, + default=MAX_QUEUE_SIZE, + help="Max Queue Size", +) +parser.add_argument("--timeout", type=float, default=TIMEOUT, help="Timeout") +parser.add_argument( + "--safety-checker", + dest="safety_checker", + action="store_true", + default=SAFETY_CHECKER, + help="Safety Checker", +) +parser.add_argument( + "--taesd", + dest="taesd", + action="store_true", + help="Use Tiny Autoencoder", +) +parser.add_argument( + "--no-taesd", + dest="taesd", + action="store_false", + help="Use Tiny Autoencoder", +) +parser.add_argument( + "--ssl-certfile", + dest="ssl_certfile", + type=str, + default=None, + help="SSL certfile", +) +parser.add_argument( + "--ssl-keyfile", + dest="ssl_keyfile", + type=str, + default=None, + help="SSL keyfile", +) +parser.add_argument( + "--debug", + action="store_true", + default=False, + help="Debug", +) +parser.add_argument( + "--acceleration", + type=str, + default=ACCELERATION, + choices=["none", "xformers", "sfast", "tensorrt"], + help="Acceleration", +) +parser.add_argument( + "--engine-dir", + dest="engine_dir", + type=str, + default=ENGINE_DIR, + help="Engine Dir", +) +parser.set_defaults(taesd=USE_TAESD) +config = Args(**vars(parser.parse_args())) +config.pretty_print() diff --git a/demo_w_camera/connection_manager.py b/demo_w_camera/connection_manager.py new file mode 100644 index 0000000..5230d68 --- /dev/null +++ b/demo_w_camera/connection_manager.py @@ -0,0 +1,116 @@ +from typing import Dict, Union +from uuid import UUID +import asyncio +from fastapi import WebSocket +from starlette.websockets import WebSocketState +import logging +from types import SimpleNamespace + +Connections = Dict[UUID, Dict[str, Union[WebSocket, asyncio.Queue]]] + + +class ServerFullException(Exception): + """Exception raised when the server is full.""" + + pass + + +class ConnectionManager: + def __init__(self): + self.active_connections: Connections = {} + + async def connect( + self, user_id: UUID, websocket: WebSocket, max_queue_size: int = 0 + ): + await websocket.accept() + user_count = self.get_user_count() + print(f"User count: {user_count}") + if max_queue_size > 0 and user_count >= max_queue_size: + print("Server is full") + await websocket.send_json({"status": "error", "message": "Server is full"}) + await websocket.close() + raise ServerFullException("Server is full") + print(f"New user connected: {user_id}") + self.active_connections[user_id] = { + "websocket": websocket, + "queue": asyncio.Queue(), + } + await websocket.send_json( + {"status": "connected", "message": "Connected"}, + ) + await websocket.send_json({"status": "wait"}) + await websocket.send_json({"status": "send_frame"}) + + def check_user(self, user_id: UUID) -> bool: + return user_id in self.active_connections + + async def update_data(self, user_id: UUID, new_data: SimpleNamespace): + user_session = self.active_connections.get(user_id) + if user_session: + queue = user_session["queue"] + while not queue.empty(): + try: + queue.get_nowait() + except asyncio.QueueEmpty: + continue + await queue.put(new_data) + + async def get_latest_data(self, user_id: UUID) -> SimpleNamespace: + user_session = self.active_connections.get(user_id) + if user_session: + queue = user_session["queue"] + try: + return await queue.get() + except asyncio.QueueEmpty: + return None + + def delete_user(self, user_id: UUID): + user_session = self.active_connections.pop(user_id, None) + if user_session: + queue = user_session["queue"] + while not queue.empty(): + try: + queue.get_nowait() + except asyncio.QueueEmpty: + continue + + def get_user_count(self) -> int: + return len(self.active_connections) + + def get_websocket(self, user_id: UUID) -> WebSocket: + user_session = self.active_connections.get(user_id) + if user_session: + websocket = user_session["websocket"] + if websocket.client_state == WebSocketState.CONNECTED: + return user_session["websocket"] + return None + + async def disconnect(self, user_id: UUID): + websocket = self.get_websocket(user_id) + if websocket: + await websocket.close() + self.delete_user(user_id) + + async def send_json(self, user_id: UUID, data: Dict): + try: + websocket = self.get_websocket(user_id) + if websocket: + await websocket.send_json(data) + except Exception as e: + logging.error(f"Error: Send json: {e}") + + async def receive_json(self, user_id: UUID) -> Dict: + try: + websocket = self.get_websocket(user_id) + if websocket: + return await websocket.receive_json() + except Exception as e: + logging.error(f"Error: Receive json: {e}") + + async def receive_bytes(self, user_id: UUID) -> bytes: + try: + websocket = self.get_websocket(user_id) + if websocket: + return await websocket.receive_bytes() + except Exception as e: + logging.error(f"Error: Receive bytes: {e}") diff --git a/demo_w_camera/main.py b/demo_w_camera/main.py new file mode 100644 index 0000000..415c1ac --- /dev/null +++ b/demo_w_camera/main.py @@ -0,0 +1,171 @@ +import asyncio +import base64 +import logging +import os +import sys +from io import BytesIO +from pathlib import Path + +import uvicorn +from config import Config +from fastapi import FastAPI +from fastapi.middleware.cors import CORSMiddleware +from fastapi.staticfiles import StaticFiles +from PIL import Image +from pydantic import BaseModel + + +sys.path.append(os.path.join(os.path.dirname(__file__), "..")) + +from utils.wrapper import StreamV2VWrapper + +base_model = "stabilityai/sd-turbo" + +logger = logging.getLogger("uvicorn") +PROJECT_DIR = Path(__file__).parent.parent + + +class PredictInputModel(BaseModel): + """ + The input model for the /predict endpoint. + """ + + prompt: str + + +class PredictResponseModel(BaseModel): + """ + The response model for the /predict endpoint. + """ + + base64_image: str + + +class UpdatePromptResponseModel(BaseModel): + """ + The response model for the /update_prompt endpoint. + """ + + prompt: str + + +class Api: + def __init__(self, config: Config) -> None: + """ + Initialize the API. + + Parameters + ---------- + config : Config + The configuration. + """ + self.config = config + self.stream = StreamV2VWrapper( + model_id_or_path=base_model, + t_index_list=[0], + mode='txt2img', + frame_buffer_size=1, + warmup=10, + acceleration=config.acceleration, + do_add_noise=True, + output_type="pil", + use_denoising_batch=False, + use_cached_attn=True, + use_feature_injection=True, + feature_injection_strength=0.8, + feature_similarity_threshold=0.98, + cache_interval=1, + cache_maxframes=4, + use_tome_cache=False, + cfg_type="none", + seed=1, + ) + + self.app = FastAPI() + self.app.add_api_route( + "/api/predict", + self._predict, + methods=["POST"], + response_model=PredictResponseModel, + ) + self.app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], + ) + self.app.mount("/", StaticFiles(directory="./frontend/dist", html=True), name="public") + + self._predict_lock = asyncio.Lock() + self._update_prompt_lock = asyncio.Lock() + + + async def _predict(self, inp: PredictInputModel) -> PredictResponseModel: + """ + Predict an image and return. + + Parameters + ---------- + inp : PredictInputModel + The input. + + Returns + ------- + PredictResponseModel + The prediction result. + """ + async with self._predict_lock: + return PredictResponseModel(base64_image=self._pil_to_base64(self.stream(prompt=inp.prompt))) + + def _pil_to_base64(self, image: Image.Image, format: str = "JPEG") -> bytes: + """ + Convert a PIL image to base64. + + Parameters + ---------- + image : Image.Image + The PIL image. + + format : str + The image format, by default "JPEG". + + Returns + ------- + bytes + The base64 image. + """ + buffered = BytesIO() + image.convert("RGB").save(buffered, format=format) + return base64.b64encode(buffered.getvalue()).decode("ascii") + + def _base64_to_pil(self, base64_image: str) -> Image.Image: + """ + Convert a base64 image to PIL. + + Parameters + ---------- + base64_image : str + The base64 image. + + Returns + ------- + Image.Image + The PIL image. + """ + if "base64," in base64_image: + base64_image = base64_image.split("base64,")[1] + return Image.open(BytesIO(base64.b64decode(base64_image))).convert("RGB") + + +if __name__ == "__main__": + from config import Config + + config = Config() + + uvicorn.run( + Api(config).app, + host=config.host, + port=config.port, + workers=config.workers, + ) diff --git a/demo_w_camera/util.py b/demo_w_camera/util.py new file mode 100644 index 0000000..059a3fb --- /dev/null +++ b/demo_w_camera/util.py @@ -0,0 +1,42 @@ +from importlib import import_module +from types import ModuleType +from typing import Dict, Any +from pydantic import BaseModel as PydanticBaseModel, Field +from PIL import Image +import io + + +def get_pipeline_class(pipeline_name: str) -> ModuleType: + try: + module = import_module(f"pipelines.{pipeline_name}") + except ModuleNotFoundError: + raise ValueError(f"Pipeline {pipeline_name} module not found") + + pipeline_class = getattr(module, "Pipeline", None) + + if pipeline_class is None: + raise ValueError(f"'Pipeline' class not found in module '{pipeline_name}'.") + + return pipeline_class + + +def bytes_to_pil(image_bytes: bytes) -> Image.Image: + image = Image.open(io.BytesIO(image_bytes)) + return image + + +def pil_to_frame(image: Image.Image) -> bytes: + frame_data = io.BytesIO() + image.save(frame_data, format="JPEG") + frame_data = frame_data.getvalue() + return ( + b"--frame\r\n" + + b"Content-Type: image/jpeg\r\n" + + f"Content-Length: {len(frame_data)}\r\n\r\n".encode() + + frame_data + + b"\r\n" + ) + + +def is_firefox(user_agent: str) -> bool: + return "Firefox" in user_agent diff --git a/demo_w_camera/vid2vid.py b/demo_w_camera/vid2vid.py new file mode 100644 index 0000000..9a7efc2 --- /dev/null +++ b/demo_w_camera/vid2vid.py @@ -0,0 +1,168 @@ +import sys +import os + +sys.path.append( + os.path.join( + os.path.dirname(__file__), + "..", + ) +) + +from utils.wrapper import StreamV2VWrapper + +import torch + +from config import Args +from pydantic import BaseModel, Field +from PIL import Image +import math + +base_model = "runwayml/stable-diffusion-v1-5" + +default_prompt = "A man is talking" + +page_content = """

StreamV2V

+

+ This demo showcases + StreamV2V + +video-to-video pipeline using + 4-step LCM LORA with a MJPEG stream server. +

+

+The base model is SD 1.5. We also build in some LORAs + for better stylization. +

+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "StreamV2V" + input_mode: str = "image" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + # negative_prompt: str = Field( + # default_negative_prompt, + # title="Negative Prompt", + # field="textarea", + # id="negative_prompt", + # ) + width: int = Field( + 512, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 512, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + params = self.InputParams() + + self.stream = StreamV2VWrapper( + model_id_or_path=base_model, + t_index_list=[30, 35, 40, 45], + frame_buffer_size=1, + width=params.width, + height=params.height, + warmup=10, + acceleration=args.acceleration, + do_add_noise=True, + output_type="pil", + use_denoising_batch=True, + use_cached_attn=True, + use_feature_injection=True, + feature_injection_strength=0.8, + feature_similarity_threshold=0.98, + cache_interval=4, + cache_maxframes=1, + use_tome_cache=True, + seed=1, + ) + self._init_lora() + self.last_prompt = default_prompt + self.stream.prepare( + prompt=default_prompt, + num_inference_steps=50, + guidance_scale=1.0, + ) + + self.lora_active = False + self.lora_trigger_words = ['pixelart', 'pixel art', 'Pixel art', 'PixArFK' + 'lowpoly', 'low poly', 'Low poly', + 'Claymation', 'claymation', + 'crayons', 'Crayons', 'crayons doodle', 'Crayons doodle', + 'sketch', 'Sketch', 'pencil drawing', 'Pencil drawing', + 'oil painting', 'Oil painting'] + + def _init_lora(self): + self.stream.stream.load_lora("./lora_weights/PixelArtRedmond15V-PixelArt-PIXARFK.safetensors", adapter_name='pixelart') + self.stream.stream.load_lora("./lora_weights/low_poly.safetensors", adapter_name='lowpoly') + self.stream.stream.load_lora("./lora_weights/Claymation.safetensors", adapter_name='claymation') + self.stream.stream.load_lora("./lora_weights/doodle.safetensors", adapter_name='crayons') + self.stream.stream.load_lora("./lora_weights/Sketch_offcolor.safetensors", adapter_name='sketch') + self.stream.stream.load_lora("./lora_weights/bichu-v0612.safetensors", adapter_name='oilpainting') + + def _activate_lora(self, prompt: str): + + if any(word in prompt for word in ['pixelart', 'pixel art', 'Pixel art', 'PixArFK']): + self.stream.stream.pipe.set_adapters(["lcm", "pixelart"], adapter_weights=[1.0, 1.0]) + print("Use LORA: pixelart in ./lora_weights/PixelArtRedmond15V-PixelArt-PIXARFK.safetensors") + elif any(word in prompt for word in ['lowpoly', 'low poly', 'Low poly']): + self.stream.stream.pipe.set_adapters(["lcm", "lowpoly"], adapter_weights=[1.0, 1.0]) + print("Use LORA: lowpoly in ./lora_weights/low_poly.safetensors") + elif any(word in prompt for word in ['Claymation', 'claymation']): + self.stream.stream.pipe.set_adapters(["lcm", "claymation"], adapter_weights=[1.0, 1.0]) + print("Use LORA: claymation in ./lora_weights/Claymation.safetensors") + elif any(word in prompt for word in ['crayons', 'Crayons', 'crayons doodle', 'Crayons doodle']): + self.stream.stream.pipe.set_adapters(["lcm", "crayons"], adapter_weights=[1.0, 1.0]) + print("Use LORA: crayons in ./lora_weights/doodle.safetensors") + elif any(word in prompt for word in ['sketch', 'Sketch', 'pencil drawing', 'Pencil drawing']): + self.stream.stream.pipe.set_adapters(["lcm", "sketch"], adapter_weights=[1.0, 1.0]) + print("Use LORA: sketch in ./lora_weights/Sketch_offcolor.safetensors") + elif any(word in prompt for word in ['oil painting', 'Oil painting']): + self.stream.stream.pipe.set_adapters(["lcm", "oilpainting"], adapter_weights=[1.0, 1.0]) + print("Use LORA: oilpainting in ./lora_weights/bichu-v0612.safetensors") + + def _deactivate_lora(self): + self.stream.stream.pipe.set_adapters("lcm") + print("Deactivate LORA, back to SD1.5") + + def _check_trigger_words(self, prompt: str): + return any(word in prompt for word in self.lora_trigger_words) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + + if self._check_trigger_words(params.prompt): + if not self.lora_active: + self._activate_lora(params.prompt) + self.lora_active = True + else: + if self.lora_active: + self._deactivate_lora() + self.lora_active = False + + image_tensor = self.stream.preprocess_image(params.image) + output_image = self.stream(image=image_tensor, prompt=params.prompt) + + return output_image