remove temp storing to disk

This commit is contained in:
harisreedhar
2026-02-20 22:31:11 +05:30
committed by henryruhs
parent a347c6081c
commit 1412820c34
+10 -53
View File
@@ -1,48 +1,11 @@
import os
import uuid
from time import time
from typing import Optional
import cv2
import numpy
from starlette.websockets import WebSocket, WebSocketDisconnect
from facefusion import ffmpeg, process_manager, session_context, session_manager, state_manager
from facefusion import session_context, session_manager, state_manager
from facefusion.apis.api_helper import get_sec_websocket_protocol
from facefusion.apis.endpoints.session import extract_access_token
from facefusion.filesystem import create_directory, remove_file
from facefusion.workflows import image_to_image
def save_image_bytes(image_bytes : bytes, temp_path : str, file_extension : str) -> str:
raw_path = os.path.join(temp_path, str(uuid.uuid4()) + file_extension)
with open(raw_path, 'wb') as raw_file:
raw_file.write(image_bytes)
return raw_path
def sanitize_upload(raw_path : str, temp_path : str, file_extension : str) -> Optional[str]:
target_path = os.path.join(temp_path, str(uuid.uuid4()) + file_extension)
process_manager.start()
sanitized = ffmpeg.sanitize_image(raw_path, target_path)
process_manager.end()
remove_file(raw_path)
if sanitized:
return target_path
return None
def process_target(target_path : str, output_path : str) -> bool:
state_manager.set_item('target_path', target_path)
state_manager.set_item('output_path', output_path)
state_manager.set_item('workflow', 'image-to-image')
error_code = image_to_image.process(time())
remove_file(target_path)
return error_code == 0
from facefusion.streamer import process_stream_frame
async def websocket_process_image(websocket : WebSocket) -> None:
@@ -56,23 +19,17 @@ async def websocket_process_image(websocket : WebSocket) -> None:
await websocket.accept(subprotocol = subprotocol)
if source_paths:
temp_path = state_manager.get_temp_path()
create_directory(temp_path)
try:
while True:
image_bytes = await websocket.receive_bytes()
raw_path = save_image_bytes(image_bytes, temp_path, '.jpg')
target_path = sanitize_upload(raw_path, temp_path, '.jpg')
target_vision_frame = cv2.imdecode(numpy.frombuffer(image_bytes, numpy.uint8), cv2.IMREAD_COLOR)
if target_path:
output_path = os.path.join(temp_path, str(uuid.uuid4()) + '.jpg')
if numpy.any(target_vision_frame):
temp_vision_frame = process_stream_frame(target_vision_frame)
success, result_bytes = cv2.imencode('.jpg', temp_vision_frame)
if process_target(target_path, output_path):
with open(output_path, 'rb') as output_file:
await websocket.send_bytes(output_file.read())
remove_file(output_path)
if success:
await websocket.send_bytes(result_bytes.tobytes())
except (WebSocketDisconnect, OSError):
pass