From e42c2727af4d5fb4bdd367b22887f9f8dfdb2d9a Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 4 May 2026 12:12:42 +0000 Subject: [PATCH] WAI-4748 --- __init__.py | 8 +- nodes/__init__.py | 2 + nodes/common.py | 84 +++++++++- nodes/video_nodes/green_screen_video_node.py | 118 +++++++++++++ .../replace_video_background_node.py | 157 ++++++++++++++++++ 5 files changed, 367 insertions(+), 2 deletions(-) create mode 100644 nodes/video_nodes/green_screen_video_node.py create mode 100644 nodes/video_nodes/replace_video_background_node.py diff --git a/__init__.py b/__init__.py index 8e08b09..e48e853 100644 --- a/__init__.py +++ b/__init__.py @@ -33,6 +33,8 @@ from .nodes import ( ShotByTextCustomCoordinatesNode, AttributionByImageNode, RemoveVideoBackgroundNode, + GreenScreenVideoNode, + ReplaceVideoBackgroundNode, VideoSolidColorBackgroundNode, VideoMaskByPromptNode, VideoMaskByKeyPointsNode, @@ -82,6 +84,8 @@ NODE_CLASS_MAPPINGS = { "GenerateStructuredPromptNodeV2": GenerateStructuredPromptNodeV2, "GenerateStructuredPromptLiteNodeV2": GenerateStructuredPromptLiteNodeV2, "RemoveVideoBackgroundNode":RemoveVideoBackgroundNode, + "GreenScreenVideoNode": GreenScreenVideoNode, + "ReplaceVideoBackgroundNode": ReplaceVideoBackgroundNode, "VideoSolidColorBackgroundNode":VideoSolidColorBackgroundNode, "VideoMaskByPromptNode":VideoMaskByPromptNode, "VideoMaskByKeyPointsNode":VideoMaskByKeyPointsNode, @@ -130,7 +134,9 @@ NODE_DISPLAY_NAME_MAPPINGS = { "RefineImageLiteNodeV2": "FIBO - Refine Image - Lite", "GenerateStructuredPromptNodeV2": "FIBO - Generate Structured Prompt", "GenerateStructuredPromptLiteNodeV2": "FIBO - Generate Structured Prompt - Lite", - "RemoveVideoBackgroundNode": "Bria Remove Video Background", + "RemoveVideoBackgroundNode": "Bria Video Remove Background", + "GreenScreenVideoNode": "Bria Video Green Screen", + "ReplaceVideoBackgroundNode": "Bria Video Replace Background", "VideoSolidColorBackgroundNode":"Bria SolidColor Background Video", "VideoMaskByPromptNode":"Bria Video Mask By Prompt", "VideoMaskByKeyPointsNode":"Bria Video Mask By Key Points", diff --git a/nodes/__init__.py b/nodes/__init__.py index 7872c77..ea1baa9 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -34,6 +34,8 @@ from .shot_by_image_manual_placement_node import ShotByImageManualPlacementNode from .shot_by_image_manual_padding_node import ShotByImageManualPaddingNode from .attribution_by_image_node import AttributionByImageNode from .video_nodes.remove_video_background_node import RemoveVideoBackgroundNode +from .video_nodes.green_screen_video_node import GreenScreenVideoNode +from .video_nodes.replace_video_background_node import ReplaceVideoBackgroundNode from .video_nodes.video_increase_resolution_node import VideoIncreaseResolutionNode from .video_nodes.video_solid_color_background_node import VideoSolidColorBackgroundNode from .video_nodes.video_erase_elements_node import VideoEraseElementsNode diff --git a/nodes/common.py b/nodes/common.py index f0a614a..c08aa0b 100644 --- a/nodes/common.py +++ b/nodes/common.py @@ -6,7 +6,8 @@ import base64 from torchvision.transforms import ToPILImage import requests import time - +import os +import uuid BRIA_COMFYUI_USER_AGENT = "bria/ComfyUI" @@ -208,3 +209,84 @@ def normalize_images_input(images): raise ValueError(f"Unsupported input type: {type(images)}") + +_EXT_TO_PIL_AND_MIME = { + ".png": ("PNG", "image/png"), + ".jpg": ("JPEG", "image/jpeg"), + ".jpeg": ("JPEG", "image/jpeg"), + ".webp": ("WEBP", "image/webp"), + ".gif": ("GIF", "image/gif"), + ".bmp": ("BMP", "image/bmp"), + ".tif": ("TIFF", "image/tiff"), + ".tiff": ("TIFF", "image/tiff"), +} + + +def _pil_format_and_mime_for_filename(file_name): + """Return (pil_format, content_type, file_name_with_ext). Uses .png only when there is no extension.""" + base = file_name.strip() if file_name else "" + if not base: + base = f"{uuid.uuid4()}_background" + root, ext = os.path.splitext(base) + ext = ext.lower() + if not ext: + ext = ".png" + base = f"{root}{ext}" + elif ext not in _EXT_TO_PIL_AND_MIME: + ext = ".png" + base = f"{root}{ext}" + pil_format, mime = _EXT_TO_PIL_AND_MIME[ext] + return pil_format, mime, base + + +def upload_pil_image_to_temp(pil_image, api_token, file_name=None): + """ + Request an anonymous presigned PUT URL, upload the image bytes, return the public temp URL. + + ``file_name`` keeps its extension for format and Content-Type; if it has no extension, ``.png`` + is appended. Matches platform POST /upload-image/anonymous/presigned-url (same pattern as video). + """ + api_url = "https://platform.prod.bria-api.com/upload-image/anonymous/presigned-url" + headers = {"Content-Type": "application/json"} + if api_token: + headers["api_token"] = api_token + # Hyphenated header survives proxies that strip underscores (e.g. nginx default). + headers["X-Api-Token"] = api_token + + pil_format, content_type, file_name = _pil_format_and_mime_for_filename(file_name or "") + + payload = { + "file_name": file_name, + "content_type": content_type, + } + + buf = io.BytesIO() + to_save = pil_image + if pil_format == "JPEG" and to_save.mode in ("RGBA", "P"): + to_save = to_save.convert("RGB") + save_kwargs = {} + if pil_format == "JPEG": + save_kwargs["quality"] = 95 + to_save.save(buf, format=pil_format, **save_kwargs) + buf.seek(0) + image_bytes = buf.read() + response = requests.post(api_url, json=payload, headers=headers) + if response.status_code != 200: + raise Exception(f"Failed to get image presigned URL: {response.status_code} {response.text}") + + response_data = response.json() + image_url = response_data.get("image_url") + upload_url = response_data.get("upload_url") + + if not image_url or not upload_url: + raise Exception(f"Invalid response from image presigned URL API: {response_data}") + + upload_response = requests.put( + upload_url, + data=image_bytes, + headers={"Content-Type": content_type}, + ) + if upload_response.status_code not in (200, 204): + raise Exception(f"Failed to upload image to S3: {upload_response.status_code}") + + return image_url \ No newline at end of file diff --git a/nodes/video_nodes/green_screen_video_node.py b/nodes/video_nodes/green_screen_video_node.py new file mode 100644 index 0000000..027a79b --- /dev/null +++ b/nodes/video_nodes/green_screen_video_node.py @@ -0,0 +1,118 @@ +import os +import uuid +import requests +from ..common import ( + poll_status_until_completed, +) +from .video_utils import upload_video_to_s3 + + +class GreenScreenVideoNode(): + """ + Applies green-screen (chroma key) background removal using the Bria API + (POST /v2/video/edit/green_screen). Output is a processed video with a solid-color background. + """ + @classmethod + def INPUT_TYPES(self): + return { + "required": { + "api_key": ("STRING", {"default": "BRIA_API_TOKEN"}), + "video_url": ("STRING", { + "default": "", + "tooltip": "Local path or publicly accessible URL of the video to process.", + }), + }, + "optional": { + "green_shade": ([ + "broadcast_green", + "chroma_green", + "blue_screen", + ], {"default": "broadcast_green"}), + "output_container_and_codec": ([ + "mp4_h264", + "mp4_h265", + "webm_vp9", + "mov_h265", + "mov_proresks", + "mkv_h264", + "mkv_h265", + "mkv_vp9", + "gif" + ], {"default": "mp4_h264"}), + "preserve_audio": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("result_video_url",) + CATEGORY = "API Nodes" + FUNCTION = "execute" + + def __init__(self): + self.api_url = "https://engine.prod.bria-api.com/v2/video/edit/green_screen" + + def execute( + self, + api_key, + video_url, + green_shade="broadcast_green", + output_container_and_codec="mp4_h264", + preserve_audio=True, + ): + if api_key.strip() == "" or api_key.strip() == "BRIA_API_TOKEN": + raise Exception("Please insert a valid API key.") + + if not video_url or not str(video_url).strip(): + raise Exception("video_url is required: provide a local path or a publicly accessible video URL.") + + if os.path.exists(video_url): + filename = f"{str(uuid.uuid4())}_{os.path.basename(video_url)}" + input_video_url = upload_video_to_s3(video_url, filename, api_key) + if not input_video_url or not ( + input_video_url.startswith("http://") or input_video_url.startswith("https://") + ): + raise Exception(f"Failed to upload video to S3. Got: {input_video_url}") + else: + input_video_url = video_url.strip() + + try: + print("Calling Bria API for video green screen...") + payload = { + "video": input_video_url, + "green_shade": green_shade, + "output_container_and_codec": output_container_and_codec, + "preserve_audio": preserve_audio, + } + + headers = { + "Content-Type": "application/json", + "api_token": f"{api_key}", + } + + response = requests.post(self.api_url, json=payload, headers=headers) + + if response.status_code == 200 or response.status_code == 202: + print("Initial video green-screen request accepted, polling for completion...") + response_dict = response.json() + + status_url = response_dict.get("status_url") + request_id = response_dict.get("request_id") + + if not status_url: + raise Exception("No status_url returned from API") + + print(f"Request ID: {request_id}, Status URL: {status_url}") + + final_response = poll_status_until_completed( + status_url, api_key, timeout=3600, check_interval=5 + ) + + result_video_url = final_response["result"]["video_url"] + + print(f"Video processing completed. Result URL: {result_video_url}") + return (result_video_url,) + + raise Exception(f"Error: API request failed with status code {response.status_code} {response.text}") + + except Exception as e: + raise Exception(f"{e}") \ No newline at end of file diff --git a/nodes/video_nodes/replace_video_background_node.py b/nodes/video_nodes/replace_video_background_node.py new file mode 100644 index 0000000..bcaf540 --- /dev/null +++ b/nodes/video_nodes/replace_video_background_node.py @@ -0,0 +1,157 @@ +import os +import uuid +import requests +from ..common import ( + normalize_images_input, + poll_status_until_completed, + upload_pil_image_to_temp +) +from .video_utils import upload_video_to_s3 + + +class ReplaceVideoBackgroundNode(): + """ + Composites a new background (image or video URL, or an IMAGE from another node) behind the + foreground video using the Bria API (POST /v2/video/edit/replace_background). + + When ``background_image`` is connected, only the first image is used (no batch); it is uploaded + via the platform anonymous image presigned URL (same pattern as video) and the resulting + ``https://temp.bria.ai/...`` URL is sent in ``background_url``. + + The background asset must match the foreground aspect ratio; otherwise the API may return + BACKGROUND_ASPECT_RATIO_MISMATCH (surfaced with foreground and background aspect ratio values). + """ + @classmethod + def INPUT_TYPES(self): + return { + "required": { + "api_key": ("STRING", {"default": "BRIA_API_TOKEN"}), + "video_url": ("STRING", { + "default": "", + "tooltip": "Local path or publicly accessible URL of the foreground video.", + }), + }, + "optional": { + "background_url": ("STRING", { + "default": "", + "tooltip": "Public HTTPS image or video URL, if not using background_image.", + }), + "background_image": ("IMAGE",), + "background_type": (["image", "video"], {"default": "image"}), + "output_container_and_codec": ([ + "mp4_h264", + "mp4_h265", + "webm_vp9", + "mov_h265", + "mov_proresks", + "mkv_h264", + "mkv_h265", + "mkv_vp9", + "gif" + ], {"default": "mp4_h264"}), + "preserve_audio": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("result_video_url",) + CATEGORY = "API Nodes" + FUNCTION = "execute" + + def __init__(self): + self.api_url = "https://engine.prod.bria-api.com/v2/video/edit/replace_background" + + @staticmethod + def _background_image_to_temp_url(background_image, api_key): + """First image only; upload to temp bucket (format from file_name extension; .png if none).""" + if background_image is None: + return None + try: + pil_images = normalize_images_input(background_image) + except (ValueError, TypeError) as e: + raise Exception(f"Invalid background_image: {e}") from e + if not pil_images: + raise Exception("background_image produced no images.") + file_name = f"{uuid.uuid4()}_background" + return upload_pil_image_to_temp(pil_images[0], api_key, file_name=file_name) + + def execute( + self, + api_key, + video_url, + background_url="", + background_image=None, + background_type="image", + output_container_and_codec="mp4_h264", + preserve_audio=True, + ): + if api_key.strip() == "" or api_key.strip() == "BRIA_API_TOKEN": + raise Exception("Please insert a valid API key.") + + if not video_url or not str(video_url).strip(): + raise Exception("video_url is required: provide a local path or a publicly accessible video URL.") + + bg_from_image = self._background_image_to_temp_url(background_image, api_key) + bg_from_url = str(background_url).strip() if background_url else "" + + if bg_from_image: + bg = bg_from_image + elif bg_from_url: + bg = bg_from_url + else: + raise Exception( + "Provide either background_image (IMAGE from Load Image, Generate Image, etc.) " + "or a non-empty background_url (HTTPS image or video URL)." + ) + + if os.path.exists(video_url): + filename = f"{str(uuid.uuid4())}_{os.path.basename(video_url)}" + input_video_url = upload_video_to_s3(video_url, filename, api_key) + if not input_video_url or not ( + input_video_url.startswith("http://") or input_video_url.startswith("https://") + ): + raise Exception(f"Failed to upload video to S3. Got: {input_video_url}") + else: + input_video_url = video_url.strip() + + try: + print("Calling Bria API for video replace background...") + payload = { + "video": input_video_url, + "background_url": bg, + "background_type": background_type, + "output_container_and_codec": output_container_and_codec, + "preserve_audio": preserve_audio, + } + + headers = { + "Content-Type": "application/json", + "api_token": f"{api_key}", + } + + response = requests.post(self.api_url, json=payload, headers=headers) + + if response.status_code == 200 or response.status_code == 202: + print("Initial video replace-background request accepted, polling for completion...") + response_dict = response.json() + + status_url = response_dict.get("status_url") + request_id = response_dict.get("request_id") + + if not status_url: + raise Exception("No status_url returned from API") + + print(f"Request ID: {request_id}, Status URL: {status_url}") + + final_response = poll_status_until_completed( + status_url, api_key, timeout=3600, check_interval=5 + ) + + result_video_url = final_response["result"]["video_url"] + + print(f"Video processing completed. Result URL: {result_video_url}") + return (result_video_url,) + + raise Exception(f"Error: API request failed with status code {response.status_code} {response.text}") + except Exception as e: + raise Exception(f"{e}")