nilor-corp-nilor-nodes/media_stream.py

import torch
import numpy as np
from PIL import Image
import requests
import io
import logging
import imageio.v2 as imageio
import mimetypes
import boto3
import json

import tempfile
import os
from .logger import logger
from .config.config import load_nilor_nodes_config

# Load shared configuration once
_CFG = load_nilor_nodes_config()

# --- Node Categories ---
category = "Nilor Nodes 👺"
subcategories = {
    "streaming": "/Streaming",
}


# --- MediaStreamInput: Universal Media Downloader ---
class MediaStreamInput:
    """
    A custom node to download an image/video from a pre-signed URL and provide it as a tensor.
    """

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "input_name": (
                    "STRING",
                    {"default": "default_input", "multiline": False},
                ),
                "format": (["image", "image_batch", "video"],),
                "presigned_download_url": (
                    "STRING",
                    {"multiline": True, "default": "<auto-filled by system>"},
                ),
            },
            "hidden": {},
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "download"
    CATEGORY = category + subcategories["streaming"]

    def download(
            self,
            presigned_download_url: str,
            format: str,
            input_name: str = "default_input",
    ):
        logger.info(
            f"ℹ️\u2009 Nilor-Nodes: MediaStreamInput: Downloading from {presigned_download_url} for input '{input_name}' with format '{format}'"
        )
        try:
            # Two-phase download for batches: manifest first, then assets
            if format == "image_batch":
                manifest_response = requests.get(presigned_download_url, timeout=60)
                manifest_response.raise_for_status()
                manifest = manifest_response.json()

                logger.info(
                    f"ℹ️\u2009 Nilor-Nodes: Processing manifest for '{manifest.get('input_name')}' with {len(manifest.get('files', []))} assets."
                )

                # Sort files by sequence number to ensure correct order
                sorted_files = sorted(
                    manifest.get("files", []), key=lambda x: x.get("sequence", 0)
                )

                # Download all assets in parallel
                asset_responses = []
                for file_info in sorted_files:
                    try:
                        resp = requests.get(file_info["presigned_url"], timeout=180)
                        resp.raise_for_status()
                        asset_responses.append(resp.content)
                    except requests.RequestException as e:
                        logger.error(
                            f"🛑\u2009 Nilor-Nodes: Failed to download asset {file_info.get('filename')}: {e}"
                        )
                        raise  # Re-raise to fail the entire process

                return self._process_image_batch(asset_responses)

            # --- Single-file download ---
            if format == "video":
                # Stream video to temp file to avoid loading entire video into RAM
                temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
                try:
                    logger.info(
                        f"ℹ️\u2009 Nilor-Nodes (MediaStreamInput): Streaming video to temp file: {temp_file.name}")
                    with requests.get(presigned_download_url, timeout=180, stream=True) as response:
                        response.raise_for_status()
                        for chunk in response.iter_content(chunk_size=8192):
                            temp_file.write(chunk)
                    temp_file.close()
                    return self._process_video(temp_file.name)
                finally:
                    # Clean up temp file
                    if os.path.exists(temp_file.name):
                        os.unlink(temp_file.name)
            else:
                # For images, load into memory (they're small)
                response = requests.get(presigned_download_url, timeout=180)
                response.raise_for_status()
                media_bytes = response.content

                if format == "image":
                    return self._process_image(media_bytes)
                else:
                    # Should not happen if UI choices are respected
                    raise ValueError(
                        f"🛑\u2009 Nilor-Nodes (MediaStreamInput): Unsupported format '{format}' for single media download."
                    )

        except requests.RequestException as e:
            logger.error(
                f"🛑\u2009 Nilor-Nodes (MediaStreamInput): Failed to download file: {e}"
            )
            return (None,)
        except Exception as e:
            logger.error(
                f"🛑\u2009 Nilor-Nodes (MediaStreamInput): Failed to process media: {e}"
            )
            return (None,)

    def _process_image_batch(self, image_bytes_list):
        logger.info(
            f"ℹ️\u2009 Nilor-Nodes (MediaStreamInput): Processing image batch with {len(image_bytes_list)} images..."
        )
        output_images = []

        for image_bytes in image_bytes_list:
            image_pil = Image.open(io.BytesIO(image_bytes))

            rgb_image_pil = image_pil.convert("RGB")
            image_tensor = torch.from_numpy(
                np.array(rgb_image_pil).astype(np.float32) / 255.0
            ).unsqueeze(0)

            output_images.append(image_tensor)

        # Concatenate along the batch dimension (dim=0)
        images_tensor = torch.cat(output_images, dim=0)

        logger.info(
            f"✅ Nilor-Nodes (MediaStreamInput): Image batch processing successful. Batch shape: {images_tensor.shape}"
        )
        return (images_tensor,)

    def _process_image(self, image_bytes):
        logger.info("ℹ️\u2009 Nilor-Nodes (MediaStreamInput): Processing as image...")
        image_pil = Image.open(io.BytesIO(image_bytes))

        # Ensure image is in RGB
        rgb_image_pil = image_pil.convert("RGB")
        image_tensor = torch.from_numpy(
            np.array(rgb_image_pil).astype(np.float32) / 255.0
        ).unsqueeze(0)

        logger.info("✅ Nilor-Nodes (MediaStreamInput): Image processing successful.")
        return (image_tensor,)

    def _process_video(self, video_path):
        logger.info(f"ℹ️\u2009 Nilor-Nodes (MediaStreamInput): Processing video from {video_path}...")

        # Open video to get metadata first
        with imageio.get_reader(video_path, format="mp4") as reader:
            # Get video metadata
            metadata = reader.get_meta_data()
            num_frames = reader.count_frames()

            if num_frames == 0:
                raise ValueError(
                    "🛑\u2009 Nilor-Nodes (MediaStreamInput): No frames could be read from the video."
                )

            # Read first frame to get dimensions
            first_frame = reader.get_data(0)
            height, width = first_frame.shape[:2]

            logger.info(
                f"ℹ️\u2009 Nilor-Nodes (MediaStreamInput): Video has {num_frames} frames at {width}x{height}"
            )

            # Pre-allocate tensor for all frames (N, H, W, 3)
            video_tensor = torch.empty((num_frames, height, width, 3), dtype=torch.float32)

            # Process first frame (already read for dimensions)
            pil_image = Image.fromarray(first_frame).convert("RGB")
            numpy_image = np.array(pil_image).astype(np.float32) / 255.0
            video_tensor[0] = torch.from_numpy(numpy_image)

            # Read remaining frames by explicit index to avoid iterator position ambiguity
            for i in range(1, num_frames):
                frame = reader.get_data(i)
                pil_image = Image.fromarray(frame).convert("RGB")
                numpy_image = np.array(pil_image).astype(np.float32) / 255.0
                video_tensor[i] = torch.from_numpy(numpy_image)

        logger.info(
            f"✅ Nilor-Nodes (MediaStreamInput): Video processing successful. Tensor shape: {video_tensor.shape}"
        )
        return (video_tensor,)


# --- MediaStreamOutput: Universal Media Uploader & SQS Notifier ---
class MediaStreamOutput:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "output_name": (
                    "STRING",
                    {"default": "default_output", "multiline": False},
                ),
                "images": ("IMAGE",),
                "format": (["png", "mp4"],),
                "framerate": ("INT", {"default": 24, "min": 1, "max": 240, "step": 1}),
                "content_id": (
                    "STRING",
                    {"default": "<auto-filled by system>", "multiline": False},
                ),
                "venue": (
                    "STRING",
                    {"default": "<auto-filled by system>", "multiline": False},
                ),
                "canvas": (
                    "STRING",
                    {"default": "<auto-filled by system>", "multiline": False},
                ),
                "scene": (
                    "STRING",
                    {"default": "<auto-filled by system>", "multiline": False},
                ),
                "presigned_upload_url": (
                    "STRING",
                    {"multiline": True, "default": "<auto-filled by system>"},
                ),
                "job_completions_queue_url": (
                    "STRING",
                    {"multiline": True, "default": "<auto-filled by system>"},
                ),
                "output_object_keys": (
                    "STRING",
                    {"multiline": False, "default": "<auto-filled by system>"},
                ),
                "job_type": (
                    "STRING",
                    {"default": "<auto-filled by system>", "multiline": False},
                ),
            },
            "hidden": {
                "prompt": "PROMPT",
                "extra_pnginfo": "EXTRA_PNGINFO",
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("uploaded_url",)
    FUNCTION = "upload_and_notify"
    OUTPUT_NODE = True
    CATEGORY = category + subcategories["streaming"]

    def upload_and_notify(
            self,
            images,
            format,
            content_id,
            venue,
            canvas,
            scene,
            presigned_upload_url,
            job_completions_queue_url,
            output_object_keys,
            framerate,
            output_name: str = "default_output",
            prompt=None,
            extra_pnginfo=None,
            job_type: str | None = None,
    ):
        if not content_id:
            raise ValueError(
                "🛑\u2009 Nilor-Nodes (MediaStreamOutput): content_id is a required input for MediaStreamOutput."
            )

        # The `output_object_keys` is received as a string representation of a dictionary.
        # We must parse it back into a dictionary.
        final_outputs_dict = {}
        try:
            # The string may use single quotes, so we replace them for valid JSON.
            final_outputs_dict = json.loads(output_object_keys.replace("'", '"'))
        except Exception as e:
            logger.error(
                f"🛑\u2009 Nilor-Nodes (MediaStreamOutput): FATAL -- Could not parse output_object_keys from string: {output_object_keys}. Error: {e}"
            )
            final_outputs_dict = {}  # Send empty dict on failure.

        # The presigned_upload_url provided to this node is specific to its output_name.
        # We don't need to re-select it. We just need to perform the upload.
        if format == "png":
            self._upload_image(images[0], presigned_upload_url)
        elif format == "mp4":
            self._upload_video(images, presigned_upload_url, framerate)

        # This node is responsible for a single output. We find its corresponding object key.
        output_key_for_this_node = final_outputs_dict.get(output_name)
        if not output_key_for_this_node:
            logging.error(
                f"🛑\u2009 Nilor-Nodes (MediaStreamOutput): FATAL -- Could not find object key for output name '{output_name}' in output_object_keys."
            )
            # Send an empty dictionary to signal failure.
            final_outputs_for_sqs = {}
        else:
            final_outputs_for_sqs = {output_name: output_key_for_this_node}

        # After upload, send the filtered dictionary of outputs to the SQS queue.
        completion_message = {
            "content_id": content_id,
            "status": "completed",
            "venue": venue,
            "canvas": canvas,
            "scene": scene,
            "outputs": final_outputs_for_sqs,
        }
        if job_type:
            completion_message["job_type"] = job_type

        try:
            # Re-initialize the client inside the execution to ensure it picks up env vars correctly.
            sqs_client = boto3.client(
                "sqs",
                endpoint_url=_CFG.worker.sqs_endpoint_url,
                aws_access_key_id=_CFG.worker.aws_access_key_id,
                aws_secret_access_key=_CFG.worker.aws_secret_access_key,
                region_name=_CFG.worker.aws_region,
            )
            logger.debug(
                f"ℹ️\u2009 Nilor-Nodes (MediaStreamOutput): Sending completion message for content {content_id} to queue: {job_completions_queue_url}"
            )
            sqs_client.send_message(
                QueueUrl=job_completions_queue_url,
                MessageBody=json.dumps(completion_message),
            )
            logger.info(
                f"✅ Nilor-Nodes (MediaStreamOutput): Completion message sent successfully for content {content_id} to queue: {job_completions_queue_url}"
            )
        except Exception as e:
            logger.error(
                f"🛑\u2009 Nilor-Nodes (MediaStreamOutput): Failed to send completion message to SQS: {e}"
            )
            raise  # Re-raise to fail the ComfyUI job

        return {"ui": {"images": []}, "result": (presigned_upload_url,)}

    def _upload_image(self, image_tensor, url):
        logger.debug(
            "ℹ️\u2009 Nilor-Nodes (MediaStreamOutput): Uploading as PNG image..."
        )
        i = 255.0 * image_tensor.cpu().numpy()
        img_pil = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))

        buffer = io.BytesIO()
        img_pil.save(buffer, format="PNG", compress_level=4)
        buffer.seek(0)

        self._perform_upload(buffer, url, "image/png")

    def _upload_video(self, image_batch_tensor, url, framerate):
        logger.info(
            f"ℹ️\u2009 Nilor-Nodes (MediaStreamOutput): Uploading as MP4 video. Frame count: {len(image_batch_tensor)}"
        )
        frames = []
        for image_tensor in image_batch_tensor:
            i = 255.0 * image_tensor.cpu().numpy()
            frame = np.clip(i, 0, 255).astype(np.uint8)
            frames.append(frame)

        buffer = io.BytesIO()
        imageio.mimwrite(buffer, frames, format="mp4", fps=framerate, quality=8)
        buffer.seek(0)

        self._perform_upload(buffer, url, "video/mp4")

    def _perform_upload(self, buffer, url, content_type):
        try:
            logger.info(
                f"ℹ️\u2009 Nilor-Nodes (MediaStreamOutput): Uploading to {url} with Content-Type: {content_type}"
            )
            headers = {"Content-Type": content_type}
            response = requests.put(
                url, data=buffer.read(), headers=headers, timeout=300
            )
            response.raise_for_status()
            logger.info("✅ Nilor-Nodes (MediaStreamOutput): Upload successful.")
        except requests.RequestException as e:
            logger.error(
                f"🛑\u2009 Nilor-Nodes (MediaStreamOutput): Failed to upload media: {e}"
            )
            raise
        except Exception as e:
            logger.error(
                f"🛑\u2009 Nilor-Nodes (MediaStreamOutput): Failed to process and upload media: {e}"
            )
            raise


# --- Node Mappings ---
NODE_CLASS_MAPPINGS = {
    "MediaStreamInput": MediaStreamInput,
    "MediaStreamOutput": MediaStreamOutput,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "MediaStreamInput": "👺 Media Stream Input (URL)",
    "MediaStreamOutput": "👺 Media Stream Output (URL)",
}