Add Luma nodes (#16)

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

comfy_api_nodes/apis/client.py

@@ -104,7 +104,7 @@ from typing import (
     TypeVar,
     Generic,
 )
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from enum import Enum
 import json
 import requests
@@ -124,6 +124,15 @@ class EmptyRequest(BaseModel):
     pass
 
 
+class UploadRequest(BaseModel):
+    filename: str = Field(..., description="Filename to upload")
+
+
+class UploadResponse(BaseModel):
+    download_url: str = Field(..., description='URL to GET uploaded file')
+    upload_url: str = Field(..., description='URL to PUT file to upload')
+
+
 class HttpMethod(str, Enum):
     GET = "GET"
     POST = "POST"

comfy_api_nodes/apis/luma_api.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Optional, Union
+
+from pydantic import BaseModel, Field, confloat
+
+
+class LumaImageModel(str, Enum):
+    photon_1 = "photon-1"
+    photon_flash_1 = "photon-flash-1"
+
+
+class LumaVideoModel(str, Enum):
+    ray_2 = "ray-2"
+    ray_flash_2 = "ray-flash-2"
+    ray_1_6 = "ray-1-6"
+
+
+class LumaAspectRatio(str, Enum):
+    ratio_1_1 = "1:1"
+    ratio_16_9 = "16:9"
+    ratio_9_16 = "9:16"
+    ratio_4_3 = "4:3"
+    ratio_3_4 = "3:4"
+    ratio_21_9 = "21:9"
+    ratio_9_21 = "9:21"
+
+
+class LumaVideoOutputResolution(str, Enum):
+    res_540p = "540p"
+    res_720p = "720p"
+    res_1080p = "1080p"
+    res_4k = "4k"
+
+
+class LumaVideoModelOutputDuration(str, Enum):
+    dur_5s = "5s"
+    dur_9s = "9s"
+
+
+class LumaGenerationType(str, Enum):
+    video = 'video'
+    image = 'image'
+
+
+class LumaState(str, Enum):
+    queued = "queued"
+    dreaming = "dreaming"
+    completed = "completed"
+    failed = "failed"
+
+
+class LumaAssets(BaseModel):
+    video: Optional[str] = Field(None, description='The URL of the video')
+    image: Optional[str] = Field(None, description='The URL of the image')
+    progress_video: Optional[str] = Field(None, description='The URL of the progress video')
+
+
+class LumaImageRef(BaseModel):
+    '''Used for image gen'''
+    url: str = Field(..., description='The URL of the image reference')
+    weight: confloat(ge=0.0, le=1.0) = Field(..., description='The weight of the image reference')
+
+
+class LumaImageReference(BaseModel):
+    '''Used for video gen'''
+    type: str = Field('image', description='Input type, defaults to image')
+    url: str = Field(..., description='The URL of the image')
+
+
+class LumaModifyImageRef(BaseModel):
+    url: str = Field(..., description='The URL of the image reference')
+    weight: confloat(ge=0.0, le=1.0) = Field(..., description='The weight of the image reference')
+
+
+class LumaCharacterRef(BaseModel):
+    identity0: LumaImageIdentity = Field(..., description='The image identity object')
+
+
+class LumaImageIdentity(BaseModel):
+    images: list[str] = Field(..., description='The URLs of the image identity')
+
+
+class LumaGenerationReference(BaseModel):
+    type: str = Field('generation', description='Input type, defaults to generation')
+    id: str = Field(..., description='The ID of the generation')
+
+
+class LumaKeyframe(BaseModel):
+    generation: Optional[LumaGenerationReference] = Field(None, description='Reference to generation')
+    image: Optional[LumaImageReference] = Field(None, description='Reference to image')
+
+
+class LumaKeyframes(BaseModel):
+    frame0: Optional[LumaKeyframe] = Field(None, description='')
+    frame1: Optional[LumaKeyframe] = Field(None, description='')
+
+
+class LumaImageGenerationRequest(BaseModel):
+    prompt: str = Field(..., description='The prompt of the generation')
+    model: LumaImageModel = Field(LumaImageModel.photon_1, description='The image model used for the generation')
+    aspect_ratio: Optional[LumaAspectRatio] = Field(LumaAspectRatio.ratio_16_9, description='The aspect ratio of the generation')
+    image_ref: Optional[list[LumaImageRef]] = Field(None, description='List of image reference objects')
+    style_ref: Optional[list[LumaImageRef]] = Field(None, description='List of style reference objects')
+    character_ref: Optional[LumaCharacterRef] = Field(None, description='The image identity object')
+    modify_image_ref: Optional[LumaModifyImageRef] = Field(None, description='The modify image reference object')
+
+
+class LumaGenerationRequest(BaseModel):
+    prompt: str = Field(..., description='The prompt of the generation')
+    model: LumaVideoModel = Field(LumaVideoModel.ray_2, description='The video model used for the generation')
+    duration: LumaVideoModelOutputDuration = Field(LumaVideoModelOutputDuration.dur_5s, description='The duration of the generation')
+    aspect_ratio: Optional[LumaAspectRatio] = Field(None, description='The aspect ratio of the generation')
+    resolution: Optional[LumaVideoOutputResolution] = Field(None, description='The resolution of the generation')
+    loop: Optional[bool] = Field(None, description='Whether to loop the video')
+    keyframes: Optional[LumaKeyframes] = Field(None, description='The keyframes of the generation')
+
+
+class LumaGeneration(BaseModel):
+    id: str = Field(..., description='The ID of the generation')
+    generation_type: LumaGenerationType = Field(..., description='Generation type, image or video')
+    state: LumaState = Field(..., description='The state of the generation')
+    failure_reason: Optional[str] = Field(None, description='The reason for the state of the generation')
+    created_at: str = Field(..., description='The date and time when the generation was created')
+    assets: Optional[LumaAssets] = Field(None, description='The assets of the generation')
+    model: str = Field(..., description='The model used for the generation')
+    request: Union[LumaGenerationRequest, LumaImageGenerationRequest] = Field(..., description="The request used for the generation")

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",
 }