Add rest of Luma node functionality (#19)

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

comfy_api_nodes/apis/luma_api.py

@@ -1,11 +1,52 @@
 from __future__ import annotations
 
+
+import torch
+
 from enum import Enum
 from typing import Optional, Union
 
 from pydantic import BaseModel, Field, confloat
 
 
+
+class LumaIO:
+    LUMA_REF = "LUMA_REF"
+
+
+class LumaReference:
+    def __init__(self, image: torch.Tensor, weight: float):
+        self.image = image
+        self.weight = weight
+
+    def create_api_model(self, download_url: str):
+        return LumaImageRef(url=download_url, weight=self.weight)
+
+class LumaReferenceChain:
+    def __init__(self, first_ref: LumaReference=None):
+        self.refs: list[LumaReference] = []
+        if first_ref:
+            self.refs.append(first_ref)
+
+    def add(self, luma_ref: LumaReference=None):
+        self.refs.append(luma_ref)
+
+    def create_api_model(self, download_urls: list[str], max_refs=4):
+        if len(self.refs) == 0:
+            return None
+        api_refs: list[LumaImageRef] = []
+        for ref, url in zip(self.refs, download_urls):
+            api_ref = LumaImageRef(url=url, weight=ref.weight)
+            api_refs.append(api_ref)
+        return api_refs
+
+    def clone(self):
+        c = LumaReferenceChain()
+        for ref in self.refs:
+            c.add(ref)
+        return c
+
+
 class LumaImageModel(str, Enum):
     photon_1 = "photon-1"
     photon_flash_1 = "photon-flash-1"
@@ -65,7 +106,7 @@ class LumaImageRef(BaseModel):
 
 class LumaImageReference(BaseModel):
     '''Used for video gen'''
-    type: str = Field('image', description='Input type, defaults to image')
+    type: Optional[str] = Field('image', description='Input type, defaults to image')
     url: str = Field(..., description='The URL of the image')
 
 
@@ -87,14 +128,9 @@ class LumaGenerationReference(BaseModel):
     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='')
+    frame0: Optional[Union[LumaImageReference, LumaGenerationReference]] = Field(None, description='')
-    frame1: Optional[LumaKeyframe] = Field(None, description='')
+    frame1: Optional[Union[LumaImageReference, LumaGenerationReference]] = Field(None, description='')
 
 
 class LumaImageGenerationRequest(BaseModel):

comfy_api_nodes/nodes_api.py

@@ -6,7 +6,7 @@ from comfy.comfy_types.node_typing import FileLocator
 from typing import Literal, Optional
 from comfy.utils import common_upscale
 from comfy.comfy_types.node_typing import IO, ComfyNodeABC, InputTypeDict
-from comfy.utils import common_upscale
+from comfy_api.input_impl.video_types import VideoFromFile
 from comfy_api_nodes.apis import (
     OpenAIImageEditRequest,
     OpenAIImageGenerationRequest,
@@ -37,6 +37,11 @@ from comfy_api_nodes.apis.luma_api import (
     LumaCharacterRef,
     LumaModifyImageRef,
     LumaImageIdentity,
+    LumaReference,
+    LumaReferenceChain,
+    LumaImageReference,
+    LumaKeyframes,
+    LumaIO,
 )
 from comfy_api_nodes.apis.client import ApiClient, ApiEndpoint, HttpMethod, SynchronousOperation, PollingOperation, EmptyRequest, UploadRequest, UploadResponse
 
@@ -959,6 +964,44 @@ class FluxProUltraImageNode(ComfyNodeABC):
         img.save(img_byte_arr, format='PNG')
         return base64.b64encode(img_byte_arr.getvalue()).decode()
 
+class LumaReferenceNode:
+    """
+    Holds an image and weight for use with Luma Generate Image node.
+    """
+    RETURN_TYPES = (LumaIO.LUMA_REF,)
+    RETURN_NAMES = ("luma_ref",)
+    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    FUNCTION = "create_luma_reference"
+    CATEGORY = "api node/Luma"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": (IO.IMAGE, {
+                    "tooltip": "Image to use as reference.",
+                }),
+                "weight": (IO.FLOAT, {
+                    "default": 1.0,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.01,
+                    "tooltip": "Weight of image reference.",
+                }),
+            },
+            "optional": {
+                "luma_ref": (LumaIO.LUMA_REF,)
+            }
+        }
+
+    def create_luma_reference(self, image: torch.Tensor, weight: float, luma_ref: LumaReferenceChain=None):
+        if luma_ref is not None:
+            luma_ref = luma_ref.clone()
+        else:
+            luma_ref = LumaReferenceChain()
+        luma_ref.add(LumaReference(image=image, weight=round(weight, 2)))
+        return (luma_ref, )
+
 class LumaImageGenerationNode:
     """
     Generates images synchronously based on prompt and aspect ratio.
@@ -989,10 +1032,23 @@ class LumaImageGenerationNode:
                     "control_after_generate": True,
                     "tooltip": "Seed to determine if node should re-run; actual results are nondeterministic regardless of seed.",
                 }),
+                "style_image_weight": (IO.FLOAT, {
+                    "default": 1.0,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.01,
+                    "tooltip": "Weight of style image. Ignored if no style_image provided.",
+                }),
             },
             "optional": {
-                "character_ref_image": (IO.IMAGE, {
+                "image_luma_ref": (LumaIO.LUMA_REF, {
-                    "tooltip": "Character reference images; can be a batch of multiple, only the first 4 images will be considered."
+                    "tooltip": "Luma Reference node connection to influence generation with input images; up to 4 images can be considered."
+                }),
+                "style_image": (IO.IMAGE, {
+                    "tooltip": "Style reference image; only 1 image will be used."
+                }),
+                "character_image": (IO.IMAGE, {
+                    "tooltip": "Character reference images; can be a batch of multiple, up to 4 images can be considered."
                 })
             },
             "hidden": {
@@ -1000,11 +1056,21 @@ class LumaImageGenerationNode:
             }
         }
 
-    def api_call(self, prompt: str, model: str, aspect_ratio: str, seed, character_ref_image: torch.Tensor=None, auth_token=None, **kwargs):
+    def api_call(self, prompt: str, model: str, aspect_ratio: str, seed, style_image_weight: float,
+                 image_luma_ref: LumaReferenceChain=None, style_image: torch.Tensor=None, character_image: torch.Tensor=None,
+                 auth_token=None, **kwargs):
+        # handle image_luma_ref
+        api_image_ref = None
+        if image_luma_ref is not None:
+            api_image_ref = self._convert_luma_refs(image_luma_ref, auth_token=auth_token)
+        # handle style_luma_ref
+        api_style_ref = None
+        if style_image is not None:
+            api_style_ref = self._convert_style_image(style_image, weight=style_image_weight, auth_token=auth_token)
         # handle character_ref images
         character_ref = None
-        if character_ref_image is not None:
+        if character_image is not None:
-            download_urls = upload_images_to_comfyapi(character_ref_image, max_images=4, auth_token=auth_token)
+            download_urls = upload_images_to_comfyapi(character_image, max_images=4, auth_token=auth_token)
             character_ref = LumaCharacterRef(identity0=LumaImageIdentity(images=download_urls))
 
         operation = SynchronousOperation(
@@ -1018,7 +1084,9 @@ class LumaImageGenerationNode:
                 prompt=prompt,
                 model=model,
                 aspect_ratio=aspect_ratio,
-                character_ref=character_ref
+                image_ref=api_image_ref,
+                style_ref=api_style_ref,
+                character_ref=character_ref,
             ),
             auth_token=auth_token
         )
@@ -1042,6 +1110,21 @@ class LumaImageGenerationNode:
         img = process_image_response(img_response)
         return (img,)
 
+    def _convert_luma_refs(self, luma_ref: LumaReferenceChain, max_refs: int, auth_token=None):
+        luma_urls = []
+        ref_count = 0
+        for ref in luma_ref.refs:
+            download_urls = upload_images_to_comfyapi(ref.image, max_images=1, auth_token=auth_token)
+            luma_urls.append(download_urls[0])
+            ref_count += 1
+            if ref_count >= max_refs:
+                break
+        return luma_ref.create_api_model(download_urls=luma_urls, max_refs=max_refs)
+
+    def _convert_style_image(self, style_image: torch.Tensor, weight: float, auth_token=None):
+        chain = LumaReferenceChain(first_ref=LumaReference(image=style_image, weight=weight))
+        return self._convert_luma_refs(chain, max_refs=1, auth_token=auth_token)
+
 class LumaImageModifyNode:
     """
     Modifies images synchronously based on prompt and aspect ratio.
@@ -1127,7 +1210,7 @@ class LumaImageModifyNode:
         img = process_image_response(img_response)
         return (img,)
 
-class LumaVideoGenerationNode:
+class LumaTextToVideoGenerationNode:
     """
     Generates videos synchronously based on prompt and output_size.
     """
@@ -1135,7 +1218,7 @@ class LumaVideoGenerationNode:
         self.output_dir = folder_paths.get_output_directory()
         self.type: Literal["output"] = "output"
 
-    RETURN_TYPES = ("IMAGE",)
+    RETURN_TYPES = (IO.VIDEO,)
     DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
     FUNCTION = "api_call"
     API_NODE = True
@@ -1158,6 +1241,9 @@ class LumaVideoGenerationNode:
                     "default": LumaVideoOutputResolution.res_540p,
                 }),
                 "duration": ([dur.value for dur in LumaVideoModelOutputDuration],),
+                "loop": (IO.BOOLEAN, {
+                    "default": False,
+                }),
                 "seed": (IO.INT, {
                     "default": 0,
                     "min": 0,
@@ -1165,7 +1251,6 @@ class LumaVideoGenerationNode:
                     "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": {
             },
@@ -1174,9 +1259,8 @@ class LumaVideoGenerationNode:
             }
         }
 
-    def api_call(self, prompt: str, model: str, aspect_ratio: str, resolution: str, duration: str, seed, filename_prefix: str,
+    def api_call(self, prompt: str, model: str, aspect_ratio: str, resolution: str, duration: str, loop: bool, seed,
                  auth_token=None, **kwargs):
-        extra_pnginfo = None
         operation = SynchronousOperation(
             endpoint=ApiEndpoint(
                 path="/proxy/luma/generations",
@@ -1190,6 +1274,7 @@ class LumaVideoGenerationNode:
                 resolution=resolution,
                 aspect_ratio=aspect_ratio,
                 duration=duration,
+                loop=loop,
             ),
             auth_token=auth_token
         )
@@ -1210,55 +1295,119 @@ class LumaVideoGenerationNode:
         response_poll = operation.execute()
 
         vid_response = requests.get(response_poll.assets.video)
-        self._save_video_locally(vid_response, filename_prefix, extra_pnginfo)
+        return (VideoFromFile(BytesIO(vid_response.content)), )
 
-        return (None,)
+class LumaImageToVideoGenerationNode:
-        #return {"ui": {"images": results, "animated": (True,)}}
+    """
+    Generates videos synchronously based on prompt, input images, and output_size.
+    """
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type: Literal["output"] = "output"
 
-    def _save_video_locally(self, response: requests.Response, filename_prefix: str, extra_pnginfo):
+    RETURN_TYPES = (IO.VIDEO,)
-        # Construct the save path
+    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
-        full_output_folder, filename, counter, subfolder, filename_prefix = (
+    FUNCTION = "api_call"
-            folder_paths.get_save_image_path(filename_prefix, self.output_dir)
+    API_NODE = True
-        )
+    CATEGORY = "api node"
-        file_basename = f"{filename}_{counter:05}_.mp4"
-        save_path = os.path.join(full_output_folder, file_basename)
 
-        video_data = response.content
+    @classmethod
-
+    def INPUT_TYPES(s):
-        # Save the video data to a file
+        return {
-        with open(save_path, "wb") as video_file:
+            "required": {
-            video_file.write(video_data)
+                "prompt": (IO.STRING, {
-
+                    "multiline": True,
-        # Add workflow metadata to the video container
+                    "default": "",
-        #if prompt is not None or extra_pnginfo is not None:
+                    "tooltip": "Prompt for the video generation",
-        if extra_pnginfo is not None:
+                }),
-            try:
+                "model": ([model.value for model in LumaVideoModel],),
-                container = av.open(save_path, mode="r+")
+                # "aspect_ratio": ([ratio.value for ratio in LumaAspectRatio], {
-                # if prompt is not None:
+                #     "default": LumaAspectRatio.ratio_16_9,
-                #     container.metadata["prompt"] = json.dumps(prompt)
+                # }),
-                if extra_pnginfo is not None:
+                "resolution": ([resolution.value for resolution in LumaVideoOutputResolution], {
-                    for x in extra_pnginfo:
+                    "default": LumaVideoOutputResolution.res_540p,
-                        container.metadata[x] = json.dumps(extra_pnginfo[x])
+                }),
-                container.close()
+                "duration": ([dur.value for dur in LumaVideoModelOutputDuration],),
-            except Exception as e:
+                "loop": (IO.BOOLEAN, {
-                logging.warning(f"Failed to add metadata to video: {e}")
+                    "default": False,
-
+                }),
-        # Create a FileLocator for the frontend to use for the preview
+                "seed": (IO.INT, {
-        results: list[FileLocator] = [
+                    "default": 0,
-            {
+                    "min": 0,
-                "filename": file_basename,
+                    "max": 0xFFFFFFFFFFFFFFFF,
-                "subfolder": subfolder,
+                    "control_after_generate": True,
-                "type": self.type,
+                    "tooltip": "Seed to determine if node should re-run; actual results are nondeterministic regardless of seed.",
+                }),
+            },
+            "optional": {
+                "first_image": (IO.IMAGE, {
+                    "tooltip": "First frame of generated video."
+                }),
+                "last_image": (IO.IMAGE, {
+                    "tooltip": "Last frame of generated video."
+                }),
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
             }
-        ]
+        }
 
-        return results
+    def api_call(self, prompt: str, model: str, resolution: str, duration: str, loop: bool, seed,
+                 first_image: torch.Tensor=None, last_image: torch.Tensor=None,
+                 auth_token=None, **kwargs):
+        if first_image is None and last_image is None:
+            raise Exception("At least one of first_image and last_image requires an input.")
+        keyframes = self._convert_to_keyframes(first_image, last_image, auth_token)
 
-    def _get_output_type(self, output_size: str):
+        operation = SynchronousOperation(
-        if output_size in [resolution.value for resolution in LumaVideoOutputResolution]:
+            endpoint=ApiEndpoint(
-            return LumaVideoOutputResolution
+                path="/proxy/luma/generations",
-        else:
+                method=HttpMethod.POST,
-            return LumaAspectRatio
+                request_model=LumaGenerationRequest,
+                response_model=LumaGeneration
+            ),
+            request=LumaGenerationRequest(
+                prompt=prompt,
+                model=model,
+                aspect_ratio=LumaAspectRatio.ratio_16_9,
+                resolution=resolution,
+                duration=duration,
+                loop=loop,
+                keyframes=keyframes
+            ),
+            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)
+        return (VideoFromFile(BytesIO(vid_response.content)), )
+
+    def _convert_to_keyframes(self, first_image: torch.Tensor=None, last_image: torch.Tensor=None, auth_token=None):
+        if first_image is None and last_image is None:
+            return None
+        frame0 = None
+        frame1 = None
+        if first_image is not None:
+            download_urls = upload_images_to_comfyapi(first_image, max_images=1, auth_token=auth_token)
+            frame0 = LumaImageReference(type='image', url=download_urls[0])
+        if last_image is not None:
+            download_urls = upload_images_to_comfyapi(last_image, max_images=1, auth_token=auth_token)
+            frame1 = LumaImageReference(type='image', url=download_urls[0])
+        return LumaKeyframes(frame0=frame0, frame1=frame1)
 
 class MinimaxTextToVideoNode:
     """
@@ -1439,7 +1588,9 @@ NODE_CLASS_MAPPINGS = {
     "FluxProUltraImageNode": FluxProUltraImageNode,
     "LumaImageNode": LumaImageGenerationNode,
     "LumaImageModifyNode": LumaImageModifyNode,
-    "LumaVideoNode": LumaVideoGenerationNode,
+    "LumaReferenceNode": LumaReferenceNode,
+    "LumaVideoNode": LumaTextToVideoGenerationNode,
+    "LumaImageToVideoNode": LumaImageToVideoGenerationNode,
     "MinimaxTextToVideoNode": MinimaxTextToVideoNode,
 }
 
@@ -1450,8 +1601,10 @@ 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",
+    "LumaImageNode": "Luma Text to Image",
-    "LumaImageModifyNode": "Luma Modify Image",
+    "LumaImageModifyNode": "Luma Image to Image",
-    "LumaVideoNode": "Luma Generate Video",
+    "LumaReferenceNode": "Luma Reference",
+    "LumaVideoNode": "Luma Text to Video",
+    "LumaImageToVideoNode": "Luma Image to Video",
     "MinimaxTextToVideoNode": "Minimax Text to Video",
 }