Add Luma nodes (#16)

Co-authored-by: Robin Huang <robin.j.huang@gmail.com>

comfy_api_nodes/nodes_api.py

@@ -20,7 +20,21 @@ from comfy_api_nodes.apis import (
     Model
 )
 from comfy_api_nodes.apis.BFLPolling import BFLStatus
-from comfy_api_nodes.apis.client import ApiEndpoint, HttpMethod, SynchronousOperation, PollingOperation, EmptyRequest
+from comfy_api_nodes.apis.luma_api import (
+    LumaImageModel,
+    LumaVideoModel,
+    LumaVideoOutputResolution,
+    LumaVideoModelOutputDuration,
+    LumaAspectRatio,
+    LumaState,
+    LumaImageGenerationRequest,
+    LumaGenerationRequest,
+    LumaGeneration,
+    LumaCharacterRef,
+    LumaModifyImageRef,
+    LumaImageIdentity,
+)
+from comfy_api_nodes.apis.client import ApiClient, ApiEndpoint, HttpMethod, SynchronousOperation, PollingOperation, EmptyRequest, UploadRequest, UploadResponse
 
 import numpy as np
 from PIL import Image
@@ -33,6 +47,7 @@ import json
 import av
 import os
 import time
+import uuid
 import folder_paths
 
 def downscale_input(image, total_pixels=1536*1024):
@@ -106,6 +121,76 @@ def validate_aspect_ratio(aspect_ratio: str, minimum_ratio: float, maximum_ratio
             raise Exception(f"Aspect ratio cannot reduce to any greater than {maximum_ratio_str} ({maximum_ratio}), but was {aspect_ratio} ({calculated_ratio}).")
     return aspect_ratio
 
+def process_image_response(response: requests.Response):
+    '''Uses content from a Response object and converts it to a torch.Tensor'''
+    image = Image.open(io.BytesIO(response.content)).convert("RGBA")
+    image_array = np.array(image).astype(np.float32) / 255.0
+    return torch.from_numpy(image_array).unsqueeze(0)
+
+def convert_image_to_bytesio(image: torch.Tensor, name: str=None, allow_alpha=True, total_pixels=2048*2048):
+    img_binary = None
+    # only care about first image, if it is a batch
+    if len(image.shape) > 3:
+        image = image[0]
+    # TODO: remove alpha if not allowed and present
+    input_tensor = image.cpu()
+    input_tensor = downscale_input(input_tensor.unsqueeze(0), total_pixels=total_pixels).squeeze()
+    image_np = (input_tensor.numpy() * 255).astype(np.uint8)
+    img = Image.fromarray(image_np)
+    img_byte_arr = io.BytesIO()
+    img.save(img_byte_arr, format='PNG')
+    img_byte_arr.seek(0)
+    img_binary = img_byte_arr
+    img_binary.name =  f"{name if name else uuid.uuid4()}.png"
+    return img_binary
+
+def upload_images_to_comfyapi(image: torch.Tensor, max_images=8, auth_token=None) -> list[str]:
+    # if batch, try to upload each file if max_images is greater than 0
+    idx_image = 0
+    download_urls: list[str] = []
+    is_batch = len(image.shape) > 3
+    batch_length = 1
+    if is_batch:
+        batch_length = image.shape[0]
+    while True:
+        curr_image = image
+        if len(image.shape) > 3:
+            curr_image = image[idx_image]
+        # get BytesIO version of image
+        img_binary = convert_image_to_bytesio(curr_image)
+        # first, request upload/download urls from comfy API
+        operation = SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path="/customers/storage",
+                method=HttpMethod.POST,
+                request_model=UploadRequest,
+                response_model=UploadResponse
+            ),
+            request=UploadRequest(
+                filename=img_binary.name
+            ),
+            auth_token=auth_token
+        )
+        response = operation.execute()
+
+        upload_response = ApiClient.upload_file(response.upload_url, img_binary)
+        # verify success
+        try:
+            upload_response.raise_for_status()
+        except requests.exceptions.HTTPError as e:
+            raise Exception(f"Could not upload one or more images: {e}")
+        # add download_url to list
+        download_urls.append(response.download_url)
+
+        idx_image += 1
+        # stop uploading additional files if done
+        if is_batch and max_images > 0:
+            if idx_image >= max_images:
+                break
+            if idx_image >= batch_length:
+                break
+    return download_urls
+
 class OpenAIDalle2(ComfyNodeABC):
     """
     Generates images synchronously via OpenAI's DALL·E 2 endpoint.
@@ -731,7 +816,7 @@ class FluxProUltraImageNode(ComfyNodeABC):
                 if result["status"] == BFLStatus.ready:
                     img_url = result["result"]["sample"]
                     img_response = requests.get(img_url)
-                    return self._process_bfl_image_response(img_response)
+                    return process_image_response(img_response)
                 elif result["status"] in [BFLStatus.request_moderated, BFLStatus.content_moderated]:
                     status = result["status"]
                     raise Exception(f"BFL API did not return an image due to: {status}.")
@@ -753,11 +838,6 @@ class FluxProUltraImageNode(ComfyNodeABC):
             else:
                 raise Exception(f"BFL API encountered an error: {response.json()}")
 
-    def _process_bfl_image_response(self, response: requests.Response):
-        image = Image.open(io.BytesIO(response.content)).convert("RGBA")
-        image_array = np.array(image).astype(np.float32) / 255.0
-        return torch.from_numpy(image_array).unsqueeze(0)
-
     def _convert_image_to_base64(self, image: torch.Tensor):
         scaled_image = downscale_input(image, total_pixels=2048*2048)
         # remove batch dimension if present
@@ -769,6 +849,307 @@ class FluxProUltraImageNode(ComfyNodeABC):
         img.save(img_byte_arr, format='PNG')
         return base64.b64encode(img_byte_arr.getvalue()).decode()
 
+class LumaImageGenerationNode:
+    """
+    Generates images synchronously based on prompt and aspect ratio.
+    """
+    RETURN_TYPES = (IO.IMAGE,)
+    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    FUNCTION = "api_call"
+    API_NODE = True
+    CATEGORY = "api node"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "prompt": (IO.STRING, {
+                    "multiline": True,
+                    "default": "",
+                    "tooltip": "Prompt for the image generation",
+                }),
+                "model": ([model.value for model in LumaImageModel],),
+                "aspect_ratio": ([ratio.value for ratio in LumaAspectRatio], {
+                    "default": LumaAspectRatio.ratio_16_9,
+                }),
+                "seed": (IO.INT, {
+                    "default": 0,
+                    "min": 0,
+                    "max": 0xFFFFFFFFFFFFFFFF,
+                    "control_after_generate": True,
+                    "tooltip": "Seed to determine if node should re-run; actual results are nondeterministic regardless of seed.",
+                }),
+            },
+            "optional": {
+                "character_ref_image": (IO.IMAGE, {
+                    "tooltip": "Character reference images; can be a batch of multiple, only the first 4 images will be considered."
+                })
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+            }
+        }
+
+    def api_call(self, prompt: str, model: str, aspect_ratio: str, seed, character_ref_image: torch.Tensor=None, auth_token=None, **kwargs):
+        # handle character_ref images
+        character_ref = None
+        if character_ref_image is not None:
+            download_urls = upload_images_to_comfyapi(character_ref_image, max_images=4, auth_token=auth_token)
+            character_ref = LumaCharacterRef(identity0=LumaImageIdentity(images=download_urls))
+
+        operation = SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path="/proxy/luma/generations/image",
+                method=HttpMethod.POST,
+                request_model=LumaImageGenerationRequest,
+                response_model=LumaGeneration
+            ),
+            request=LumaImageGenerationRequest(
+                prompt=prompt,
+                model=model,
+                aspect_ratio=aspect_ratio,
+                character_ref=character_ref
+            ),
+            auth_token=auth_token
+        )
+        response_api: LumaGeneration = operation.execute()
+
+        operation = PollingOperation(
+            poll_endpoint=ApiEndpoint(
+                path=f"/proxy/luma/generations/{response_api.id}",
+                method=HttpMethod.GET,
+                request_model=EmptyRequest,
+                response_model=LumaGeneration,
+            ),
+            completed_statuses=[LumaState.completed],
+            failed_statuses=[LumaState.failed],
+            status_extractor=lambda x: x.state,
+            auth_token=auth_token,
+        )
+        response_poll = operation.execute()
+
+        img_response = requests.get(response_poll.assets.image)
+        img = process_image_response(img_response)
+        return (img,)
+
+class LumaImageModifyNode:
+    """
+    Modifies images synchronously based on prompt and aspect ratio.
+    """
+    RETURN_TYPES = (IO.IMAGE,)
+    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    FUNCTION = "api_call"
+    API_NODE = True
+    CATEGORY = "api node"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": (IO.IMAGE,),
+                "prompt": (IO.STRING, {
+                    "multiline": True,
+                    "default": "",
+                    "tooltip": "Prompt for the image generation",
+                }),
+                "image_weight": (IO.FLOAT, {
+                    "default": 1.0,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.01,
+                    "tooltip": "Weight of the image; the closer to 0.0, the less the image will be modified."
+                }),
+                "model": ([model.value for model in LumaImageModel],),
+                "seed": (IO.INT, {
+                    "default": 0,
+                    "min": 0,
+                    "max": 0xFFFFFFFFFFFFFFFF,
+                    "control_after_generate": True,
+                    "tooltip": "Seed to determine if node should re-run; actual results are nondeterministic regardless of seed.",
+                }),
+            },
+            "optional": {
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+            }
+        }
+
+    def api_call(self, prompt: str, model: str, image: torch.Tensor, image_weight: float, seed, auth_token=None, **kwargs):
+        # first, upload image
+        download_urls = upload_images_to_comfyapi(image, max_images=1, auth_token=auth_token)
+        image_url = download_urls[0]
+        # next, make Luma call with download url provided
+        operation = SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path="/proxy/luma/generations/image",
+                method=HttpMethod.POST,
+                request_model=LumaImageGenerationRequest,
+                response_model=LumaGeneration
+            ),
+            request=LumaImageGenerationRequest(
+                prompt=prompt,
+                model=model,
+                modify_image_ref=LumaModifyImageRef(
+                    url=image_url,
+                    weight=round(image_weight, 2)
+                ),
+            ),
+            auth_token=auth_token
+        )
+        response_api: LumaGeneration = operation.execute()
+
+        operation = PollingOperation(
+            poll_endpoint=ApiEndpoint(
+                path=f"/proxy/luma/generations/{response_api.id}",
+                method=HttpMethod.GET,
+                request_model=EmptyRequest,
+                response_model=LumaGeneration,
+            ),
+            completed_statuses=[LumaState.completed],
+            failed_statuses=[LumaState.failed],
+            status_extractor=lambda x: x.state,
+            auth_token=auth_token,
+        )
+        response_poll = operation.execute()
+
+        img_response = requests.get(response_poll.assets.image)
+        img = process_image_response(img_response)
+        return (img,)
+
+class LumaVideoGenerationNode:
+    """
+    Generates videos synchronously based on prompt and output_size.
+    """
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type: Literal["output"] = "output"
+
+    RETURN_TYPES = ("IMAGE",)
+    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    FUNCTION = "api_call"
+    API_NODE = True
+    CATEGORY = "api node"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "prompt": (IO.STRING, {
+                    "multiline": True,
+                    "default": "",
+                    "tooltip": "Prompt for the video generation",
+                }),
+                "model": ([model.value for model in LumaVideoModel],),
+                "aspect_ratio": ([ratio.value for ratio in LumaAspectRatio], {
+                    "default": LumaAspectRatio.ratio_16_9,
+                }),
+                "resolution": ([resolution.value for resolution in LumaVideoOutputResolution], {
+                    "default": LumaVideoOutputResolution.res_540p,
+                }),
+                "duration": ([dur.value for dur in LumaVideoModelOutputDuration],),
+                "seed": (IO.INT, {
+                    "default": 0,
+                    "min": 0,
+                    "max": 0xFFFFFFFFFFFFFFFF,
+                    "control_after_generate": True,
+                    "tooltip": "Seed to determine if node should re-run; actual results are nondeterministic regardless of seed.",
+                }),
+                "filename_prefix": ("STRING", {"default": "ComfyUI"}),
+            },
+            "optional": {
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+            }
+        }
+
+    def api_call(self, prompt: str, model: str, aspect_ratio: str, resolution: str, duration: str, seed, filename_prefix: str,
+                 auth_token=None, **kwargs):
+        extra_pnginfo = None
+        operation = SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path="/proxy/luma/generations",
+                method=HttpMethod.POST,
+                request_model=LumaGenerationRequest,
+                response_model=LumaGeneration
+            ),
+            request=LumaGenerationRequest(
+                prompt=prompt,
+                model=model,
+                resolution=resolution,
+                aspect_ratio=aspect_ratio,
+                duration=duration,
+            ),
+            auth_token=auth_token
+        )
+        response_api: LumaGeneration = operation.execute()
+
+        operation = PollingOperation(
+            poll_endpoint=ApiEndpoint(
+                path=f"/proxy/luma/generations/{response_api.id}",
+                method=HttpMethod.GET,
+                request_model=EmptyRequest,
+                response_model=LumaGeneration,
+            ),
+            completed_statuses=[LumaState.completed],
+            failed_statuses=[LumaState.failed],
+            status_extractor=lambda x: x.state,
+            auth_token=auth_token,
+        )
+        response_poll = operation.execute()
+
+        vid_response = requests.get(response_poll.assets.video)
+        self._save_video_locally(vid_response, filename_prefix, extra_pnginfo)
+
+        return (None,)
+        #return {"ui": {"images": results, "animated": (True,)}}
+
+    def _save_video_locally(self, response: requests.Response, filename_prefix: str, extra_pnginfo):
+        # Construct the save path
+        full_output_folder, filename, counter, subfolder, filename_prefix = (
+            folder_paths.get_save_image_path(filename_prefix, self.output_dir)
+        )
+        file_basename = f"{filename}_{counter:05}_.mp4"
+        save_path = os.path.join(full_output_folder, file_basename)
+
+        video_data = response.content
+
+        # Save the video data to a file
+        with open(save_path, "wb") as video_file:
+            video_file.write(video_data)
+
+        # Add workflow metadata to the video container
+        #if prompt is not None or extra_pnginfo is not None:
+        if extra_pnginfo is not None:
+            try:
+                container = av.open(save_path, mode="r+")
+                # if prompt is not None:
+                #     container.metadata["prompt"] = json.dumps(prompt)
+                if extra_pnginfo is not None:
+                    for x in extra_pnginfo:
+                        container.metadata[x] = json.dumps(extra_pnginfo[x])
+                container.close()
+            except Exception as e:
+                logging.warning(f"Failed to add metadata to video: {e}")
+
+        # Create a FileLocator for the frontend to use for the preview
+        results: list[FileLocator] = [
+            {
+                "filename": file_basename,
+                "subfolder": subfolder,
+                "type": self.type,
+            }
+        ]
+
+        return results
+
+    def _get_output_type(self, output_size: str):
+        if output_size in [resolution.value for resolution in LumaVideoOutputResolution]:
+            return LumaVideoOutputResolution
+        else:
+            return LumaAspectRatio
+
 class MinimaxTextToVideoNode:
     """
     Generates videos synchronously based on a prompt, and optional parameters using Minimax's API.
@@ -946,6 +1327,9 @@ NODE_CLASS_MAPPINGS = {
     "OpenAIGPTImage1": OpenAIGPTImage1,
     "IdeogramTextToImage": IdeogramTextToImage,
     "FluxProUltraImageNode": FluxProUltraImageNode,
+    "LumaImageNode": LumaImageGenerationNode,
+    "LumaImageModifyNode": LumaImageModifyNode,
+    "LumaVideoNode": LumaVideoGenerationNode,
     "MinimaxTextToVideoNode": MinimaxTextToVideoNode,
 }
 
@@ -956,5 +1340,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "OpenAIGPTImage1": "OpenAI GPT Image 1",
     "IdeogramTextToImage": "Ideogram Text to Image",
     "FluxProUltraImageNode": "Flux 1.1 [pro] Ultra Image",
+    "LumaImageNode": "Luma Generate Image",
+    "LumaImageModifyNode": "Luma Modify Image",
+    "LumaVideoNode": "Luma Generate Video",
     "MinimaxTextToVideoNode": "Minimax Text to Video",
 }