nodes/smthemex-ComfyUI_Streamv2v_Plus
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add files via upload

Browse Source 
add
This commit is contained in:
smthemex
2024-06-23 11:33:43 +08:00
committed by GitHub
parent b73216ec6e
commit f39dc19b89
5 changed files with 606 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
demo_w_camera/config.py
+109
View File
@@ -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()
demo_w_camera/connection_manager.py
+116
View File
@@ -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}")
demo_w_camera/main.py
+171
View File
@@ -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,
)
demo_w_camera/util.py
+42
View File
@@ -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
demo_w_camera/vid2vid.py
+168
View File
@@ -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 = """<h1 class="text-3xl font-bold">StreamV2V</h1>
<p class="text-sm">
This demo showcases
<a
href="https://jeff-liangf.github.io/projects/streamv2v/"
target="_blank"
class="text-blue-500 underline hover:no-underline">StreamV2V
</a>
video-to-video pipeline using
<a
href="https://huggingface.co/latent-consistency/lcm-lora-sdv1-5"
target="_blank"
class="text-blue-500 underline hover:no-underline">4-step LCM LORA</a
> with a MJPEG stream server.
</p>
<p class="text-sm">
The base model is <a
href="https://huggingface.co/runwayml/stable-diffusion-v1-5"
target="_blank"
class="text-blue-500 underline hover:no-underline">SD 1.5</a
>. We also build in <a
href="https://github.com/Jeff-LiangF/streamv2v/tree/main/demo_w_camera#download-lora-weights-for-better-stylization"
target="_blank"
class="text-blue-500 underline hover:no-underline">some LORAs
</a> for better stylization.
</p>
"""
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