harden: load training-dataset shards with weights_only=True (#14543)

LoadTrainingDataset was the only torch.load call in the codebase without
weights_only=True; comfy/utils.py and comfy/sd1_clip.py already pass it.
Recent PyTorch defaults to weights_only=True, so this is defense-in-depth
for installs pinned to older PyTorch. Verified a typical shard (latents +
standard conditioning) round-trips cleanly under weights_only=True.

comfy_api_nodes/apis/luma.py

@@ -10,7 +10,6 @@ from pydantic import BaseModel, Field, confloat
 class LumaIO:
     LUMA_REF = "LUMA_REF"
     LUMA_CONCEPTS = "LUMA_CONCEPTS"
-    LUMA_RAY32_KEYFRAME = "LUMA_RAY32_KEYFRAME"
 
 
 class LumaReference:
@@ -21,14 +20,13 @@ class LumaReference:
     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):
+    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):
+    def add(self, luma_ref: LumaReference=None):
         self.refs.append(luma_ref)
 
     def create_api_model(self, download_urls: list[str], max_refs=4):
@@ -126,7 +124,7 @@ def get_luma_concepts(include_none=False):
         "pull_out",
         "aerial",
         "crane_up",
-        "eye_level",
+        "eye_level"
     ]
 
 
@@ -164,8 +162,8 @@ class LumaVideoModelOutputDuration(str, Enum):
 
 
 class LumaGenerationType(str, Enum):
-    video = "video"
-    image = "image"
+    video = 'video'
+    image = 'image'
 
 
 class LumaState(str, Enum):
@@ -176,109 +174,86 @@ class LumaState(str, Enum):
 
 
 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")
+    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")
+    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: Optional[str] = Field("image", description="Input type, defaults to image")
-    url: str = Field(..., description="The URL of the image")
+    type: Optional[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")
+    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")
+    identity0: LumaImageIdentity = Field(..., description='The image identity object')
 
 
 class LumaImageIdentity(BaseModel):
-    images: list[str] = Field(..., description="The URLs of the image identity")
+    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")
+    type: str = Field('generation', description='Input type, defaults to generation')
+    id: str = Field(..., description='The ID of the generation')
 
 
 class LumaKeyframes(BaseModel):
-    frame0: Optional[Union[LumaImageReference, LumaGenerationReference]] = Field(None, description="")
-    frame1: Optional[Union[LumaImageReference, LumaGenerationReference]] = Field(None, description="")
+    frame0: Optional[Union[LumaImageReference, LumaGenerationReference]] = Field(None, description='')
+    frame1: Optional[Union[LumaImageReference, LumaGenerationReference]] = Field(None, description='')
 
 
 class LumaConceptObject(BaseModel):
-    key: str = Field(..., description="Camera Concept name")
+    key: str = Field(..., description='Camera Concept name')
 
 
 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)
-    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")
+    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: Optional[LumaVideoModelOutputDuration] = Field(None, 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")
-    concepts: Optional[list[LumaConceptObject]] = Field(None, description="Camera Concepts to apply to generation")
+    prompt: str = Field(..., description='The prompt of the generation')
+    model: LumaVideoModel = Field(LumaVideoModel.ray_2, description='The video model used for the generation')
+    duration: Optional[LumaVideoModelOutputDuration] = Field(None, 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')
+    concepts: Optional[list[LumaConceptObject]] = Field(None, description='Camera Concepts to apply to 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(...)
+    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")
 
 
 class Luma2ImageRef(BaseModel):
     url: str | None = None
     data: str | None = None
     media_type: str | None = None
-    generation_id: str | None = Field(None, description="reference a prior generation (extend / source reuse)")
-
-
-class Luma2VideoEdit(BaseModel):
-    """Edit controls for Ray 3.2 ``video_edit`` generations."""
-
-    auto_controls: bool | None = Field(None, description="derive a conditioning schedule from the source (recommended)")
-    strength: str | None = Field(None, description="'adhere_1' .. 'reimagine_3'; constrained by IO.Combo")
-
-
-class Luma2VideoOptions(BaseModel):
-    """Ray 3.2 ``video`` output settings (text / image / keyframe / edit / extend)."""
-
-    resolution: str | None = Field(None, description="360p | 540p | 720p | 1080p")
-    duration: str | None = Field(None, description="5s | 10s")
-    loop: bool | None = Field(None)
-    start_frame: Luma2ImageRef | None = Field(None)
-    end_frame: Luma2ImageRef | None = Field(None)
-    keyframes: list[Luma2ImageRef] | None = Field(None)
-    keyframe_indexes: list[int] | None = Field(None)
-    edit: Luma2VideoEdit | None = Field(None)
 
 
 class Luma2GenerationRequest(BaseModel):
@@ -291,7 +266,6 @@ class Luma2GenerationRequest(BaseModel):
     web_search: bool | None = None
     image_ref: list[Luma2ImageRef] | None = None
     source: Luma2ImageRef | None = None
-    video: Luma2VideoOptions | None = Field(None)
 
 
 class Luma2Generation(BaseModel):
@@ -303,31 +277,3 @@ class Luma2Generation(BaseModel):
     output: list[LumaImageReference] | None = None
     failure_reason: str | None = None
     failure_code: str | None = None
-
-
-# --- Ray 3.2 multi-keyframe chain ---
-
-LUMA_KEYFRAME_MODE_FRACTION = "fraction"  # value in [0.0, 1.0] of the output video duration
-LUMA_KEYFRAME_MODE_SECONDS = "seconds"  # absolute time, in seconds, from the start of the output
-
-
-class LumaRay32KeyframeItem:
-    """One guide image anchored at a position on the Ray 3.2 output timeline."""
-
-    def __init__(self, image: torch.Tensor, mode: str, value: float):
-        self.image = image
-        self.mode = mode  # LUMA_KEYFRAME_MODE_FRACTION | LUMA_KEYFRAME_MODE_SECONDS
-        self.value = value
-
-
-class LumaRay32KeyframeChain:
-    def __init__(self):
-        self.items: list[LumaRay32KeyframeItem] = []
-
-    def add(self, item: LumaRay32KeyframeItem) -> None:
-        self.items.append(item)
-
-    def clone(self) -> "LumaRay32KeyframeChain":
-        c = LumaRay32KeyframeChain()
-        c.items = list(self.items)
-        return c

comfy_api_nodes/nodes_luma.py

@@ -3,13 +3,9 @@ from typing_extensions import override
 
 from comfy_api.latest import IO, ComfyExtension, Input
 from comfy_api_nodes.apis.luma import (
-    LUMA_KEYFRAME_MODE_FRACTION,
-    LUMA_KEYFRAME_MODE_SECONDS,
     Luma2Generation,
     Luma2GenerationRequest,
     Luma2ImageRef,
-    Luma2VideoEdit,
-    Luma2VideoOptions,
     LumaAspectRatio,
     LumaCharacterRef,
     LumaConceptChain,
@@ -22,8 +18,6 @@ from comfy_api_nodes.apis.luma import (
     LumaIO,
     LumaKeyframes,
     LumaModifyImageRef,
-    LumaRay32KeyframeChain,
-    LumaRay32KeyframeItem,
     LumaReference,
     LumaReferenceChain,
     LumaVideoModel,
@@ -39,7 +33,6 @@ from comfy_api_nodes.util import (
     sync_op,
     upload_image_to_comfyapi,
     upload_images_to_comfyapi,
-    upload_video_to_comfyapi,
     validate_string,
 )
 
@@ -699,10 +692,7 @@ async def _luma2_upload_image_refs(
 async def _luma2_submit_and_poll(
     cls: type[IO.ComfyNode],
     request: Luma2GenerationRequest,
-    *,
-    estimated_duration: int | None = None,
-) -> Luma2Generation:
-    """Submit a Luma Agents generation and poll until done; returns the completed generation."""
+) -> Input.Image:
     initial = await sync_op(
         cls,
         ApiEndpoint(path="/proxy/luma_2/generations", method="POST"),
@@ -710,21 +700,21 @@ async def _luma2_submit_and_poll(
         data=request,
     )
     if not initial.id:
-        raise RuntimeError("Luma API did not return a generation id.")
+        raise RuntimeError("Luma 2 API did not return a generation id.")
     final = await poll_op(
         cls,
         ApiEndpoint(path=f"/proxy/luma_2/generations/{initial.id}", method="GET"),
         response_model=Luma2Generation,
         status_extractor=lambda r: r.state,
         progress_extractor=lambda r: None,
-        estimated_duration=estimated_duration,
     )
-    if not final.output or not final.output[0].url:
+    if not final.output:
         msg = final.failure_reason or "no output returned"
-        if final.failure_code:
-            msg = f"{msg} [{final.failure_code}]"
-        raise RuntimeError(f"Luma generation failed: {msg}")
-    return final
+        raise RuntimeError(f"Luma 2 generation failed: {msg}")
+    url = final.output[0].url
+    if not url:
+        raise RuntimeError("Luma 2 generation completed without an output URL.")
+    return await download_url_to_image_tensor(url)
 
 
 class LumaImageNode(IO.ComfyNode):
@@ -853,8 +843,7 @@ class LumaImageNode(IO.ComfyNode):
             web_search=model["web_search"],
             image_ref=await _luma2_upload_image_refs(cls, model.get("image_ref"), max_count=9),
         )
-        final = await _luma2_submit_and_poll(cls, request)
-        return IO.NodeOutput(await download_url_to_image_tensor(final.output[0].url))
+        return IO.NodeOutput(await _luma2_submit_and_poll(cls, request))
 
 
 class LumaImageEditNode(IO.ComfyNode):
@@ -940,533 +929,7 @@ class LumaImageEditNode(IO.ComfyNode):
             web_search=model["web_search"],
             image_ref=await _luma2_upload_image_refs(cls, model.get("image_ref"), max_count=8),
         )
-        final = await _luma2_submit_and_poll(cls, request)
-        return IO.NodeOutput(await download_url_to_image_tensor(final.output[0].url))
-
-
-_BADGE_RAY32_VIDEO = IO.PriceBadge(
-    depends_on=IO.PriceBadgeDepends(widgets=["resolution", "duration"]),
-    expr="""
-    (
-      $p := {
-        "360p": {"5s": 0.06, "10s": 0.18},
-        "540p": {"5s": 0.15, "10s": 0.45},
-        "720p": {"5s": 0.3, "10s": 0.9},
-        "1080p": {"5s": 1.2, "10s": 3.6}
-      };
-      {"type": "usd", "usd": $lookup($lookup($p, widgets.resolution), widgets.duration)}
-    )
-    """,
-)
-
-_BADGE_RAY32_VIDEO_5S = IO.PriceBadge(
-    depends_on=IO.PriceBadgeDepends(widgets=["resolution"]),
-    expr="""
-    (
-      $p := {"360p": 0.06, "540p": 0.15, "720p": 0.3, "1080p": 1.2};
-      {"type": "usd", "usd": $lookup($p, widgets.resolution)}
-    )
-    """,
-)
-
-_BADGE_RAY32_EDIT = IO.PriceBadge(
-    depends_on=IO.PriceBadgeDepends(widgets=["resolution"]),
-    expr="""
-    (
-      $p := {
-        "360p": {"min": 0.54, "max": 1.08},
-        "540p": {"min": 0.72, "max": 1.44},
-        "720p": {"min": 1.08, "max": 2.16},
-        "1080p": {"min": 2.16, "max": 4.32}
-      };
-      $r := $lookup($p, widgets.resolution);
-      {"type": "range_usd", "min_usd": $r.min, "max_usd": $r.max, "format": {"note": "(by source length)"}}
-    )
-    """,
-)
-
-_BADGE_RAY32_REFRAME = IO.PriceBadge(
-    depends_on=IO.PriceBadgeDepends(widgets=["resolution"]),
-    expr="""
-    (
-      $p := {"360p": 0.03, "540p": 0.06, "720p": 0.12, "1080p": 0.36};
-      {"type": "usd", "usd": $lookup($p, widgets.resolution), "format": {"suffix": "/second"}}
-    )
-    """,
-)
-
-
-def _ray32_seed_input() -> IO.Input:
-    return IO.Int.Input(
-        "seed",
-        default=0,
-        min=0,
-        max=0xFFFFFFFFFFFFFFFF,
-        control_after_generate=True,
-        tooltip="Seed to determine if node should re-run; results are nondeterministic regardless of seed.",
-    )
-
-
-async def _ray32_generate(cls: type[IO.ComfyNode], request: Luma2GenerationRequest) -> IO.NodeOutput:
-    """Run a ray-3.2 generation and return (video, generation_id)."""
-    final = await _luma2_submit_and_poll(cls, request, estimated_duration=120)
-    video = await download_url_to_video_output(final.output[0].url)
-    return IO.NodeOutput(video, final.id or "")
-
-
-class LumaRay32TextToVideoNode(IO.ComfyNode):
-    @classmethod
-    def define_schema(cls) -> IO.Schema:
-        return IO.Schema(
-            node_id="LumaRay32TextToVideoNode",
-            display_name="Luma Ray 3.2 Text to Video",
-            category="partner/video/Luma",
-            description="Generate a video from a text prompt using Luma's Ray 3.2 model.",
-            inputs=[
-                IO.String.Input("prompt", multiline=True, default="", tooltip="Text prompt for the video generation."),
-                IO.Combo.Input("aspect_ratio", options=["16:9", "9:16", "1:1", "4:3", "3:4", "21:9"]),
-                IO.Combo.Input("resolution", options=["360p", "540p", "720p", "1080p"], default="720p"),
-                IO.Combo.Input("duration", options=["5s", "10s"]),
-                IO.Boolean.Input(
-                    "loop",
-                    default=False,
-                    tooltip="Make the video loop seamlessly. Only available with 5s duration.",
-                ),
-                _ray32_seed_input(),
-            ],
-            outputs=[IO.Video.Output(), IO.String.Output(display_name="generation_id")],
-            hidden=[
-                IO.Hidden.auth_token_comfy_org,
-                IO.Hidden.api_key_comfy_org,
-                IO.Hidden.unique_id,
-            ],
-            is_api_node=True,
-            price_badge=_BADGE_RAY32_VIDEO,
-        )
-
-    @classmethod
-    async def execute(
-        cls, prompt: str, aspect_ratio: str, resolution: str, duration: str, loop: bool, seed: int
-    ) -> IO.NodeOutput:
-        validate_string(prompt, strip_whitespace=True, min_length=1, max_length=6000)
-        if loop and duration == "10s":
-            raise ValueError("Looping is only available with 5s duration on Ray 3.2.")
-        request = Luma2GenerationRequest(
-            prompt=prompt,
-            model="ray-3.2",
-            type="video",
-            aspect_ratio=aspect_ratio,
-            video=Luma2VideoOptions(resolution=resolution, duration=duration, loop=loop or None),
-        )
-        return await _ray32_generate(cls, request)
-
-
-class LumaRay32ImageToVideoNode(IO.ComfyNode):
-    @classmethod
-    def define_schema(cls) -> IO.Schema:
-        return IO.Schema(
-            node_id="LumaRay32ImageToVideoNode",
-            display_name="Luma Ray 3.2 Image to Video",
-            category="partner/video/Luma",
-            description="Generate a video from a start and/or end frame using Luma's Ray 3.2 model. "
-            "Image-anchored generations are always 5 seconds.",
-            inputs=[
-                IO.String.Input("prompt", multiline=True, default="", tooltip="Text prompt for the video generation."),
-                IO.Combo.Input("resolution", options=["360p", "540p", "720p", "1080p"], default="720p"),
-                IO.Boolean.Input(
-                    "loop",
-                    default=False,
-                    tooltip="Make the video loop seamlessly. Not available when an end_frame is set.",
-                ),
-                _ray32_seed_input(),
-                IO.Image.Input("start_frame", optional=True, tooltip="First frame of the generated video."),
-                IO.Image.Input("end_frame", optional=True, tooltip="Last frame of the generated video."),
-            ],
-            outputs=[IO.Video.Output(), IO.String.Output(display_name="generation_id")],
-            hidden=[
-                IO.Hidden.auth_token_comfy_org,
-                IO.Hidden.api_key_comfy_org,
-                IO.Hidden.unique_id,
-            ],
-            is_api_node=True,
-            price_badge=_BADGE_RAY32_VIDEO_5S,
-        )
-
-    @classmethod
-    async def execute(
-        cls,
-        prompt: str,
-        resolution: str,
-        loop: bool,
-        seed: int,
-        start_frame: torch.Tensor | None = None,
-        end_frame: torch.Tensor | None = None,
-    ) -> IO.NodeOutput:
-        validate_string(prompt, strip_whitespace=True, min_length=1, max_length=6000)
-        if start_frame is None and end_frame is None:
-            raise ValueError("Provide at least one of start_frame / end_frame.")
-        if loop and end_frame is not None:
-            raise ValueError("Looping is not available when an end_frame is set.")
-        video = Luma2VideoOptions(resolution=resolution, duration="5s", loop=loop or None)
-        if start_frame is not None:
-            url = await upload_image_to_comfyapi(cls, start_frame, mime_type="image/png")
-            video.start_frame = Luma2ImageRef(url=url)
-        if end_frame is not None:
-            url = await upload_image_to_comfyapi(cls, end_frame, mime_type="image/png")
-            video.end_frame = Luma2ImageRef(url=url)
-        request = Luma2GenerationRequest(prompt=prompt, model="ray-3.2", type="video", video=video)
-        return await _ray32_generate(cls, request)
-
-
-class LumaRay32KeyframeNode(IO.ComfyNode):
-    @classmethod
-    def define_schema(cls) -> IO.Schema:
-        return IO.Schema(
-            node_id="LumaRay32KeyframeNode",
-            display_name="Luma Ray 3.2 Keyframe",
-            category="partner/video/Luma",
-            description="Anchor a guide image to a position on the Ray 3.2 output video timeline. Connect this to "
-            "the 'keyframes' input of the Luma Ray 3.2 Keyframes to Video node; chain several together via the "
-            "optional 'keyframes' input below.",
-            inputs=[
-                IO.Image.Input("image", tooltip="Guide image to place at the chosen moment of the output video."),
-                IO.DynamicCombo.Input(
-                    "position",
-                    options=[
-                        IO.DynamicCombo.Option(
-                            "Fraction of duration (0.0-1.0)",
-                            [
-                                IO.Float.Input(
-                                    "fraction",
-                                    default=0.0,
-                                    min=0.0,
-                                    max=1.0,
-                                    step=0.01,
-                                    display_mode=IO.NumberDisplay.number,
-                                    tooltip="Where in the output video this image applies " "(0.0 = start, 1.0 = end).",
-                                ),
-                            ],
-                        ),
-                        IO.DynamicCombo.Option(
-                            "Absolute time (seconds)",
-                            [
-                                IO.Float.Input(
-                                    "seconds",
-                                    default=0.0,
-                                    min=0.0,
-                                    max=10.0,
-                                    step=0.1,
-                                    display_mode=IO.NumberDisplay.number,
-                                    tooltip="Time in seconds from the start of the output video where this "
-                                    "image applies.",
-                                ),
-                            ],
-                        ),
-                    ],
-                    tooltip="How to place this image on the output video's timeline.",
-                ),
-                IO.Custom(LumaIO.LUMA_RAY32_KEYFRAME).Input(
-                    "keyframes",
-                    optional=True,
-                    tooltip="Optional earlier keyframes to chain with this one.",
-                ),
-            ],
-            outputs=[IO.Custom(LumaIO.LUMA_RAY32_KEYFRAME).Output(display_name="keyframes")],
-        )
-
-    @classmethod
-    def execute(
-        cls,
-        image: torch.Tensor,
-        position: dict,
-        keyframes: LumaRay32KeyframeChain | None = None,
-    ) -> IO.NodeOutput:
-        chain = keyframes.clone() if keyframes is not None else LumaRay32KeyframeChain()
-        if position["position"] == "Absolute time (seconds)":
-            mode, value = LUMA_KEYFRAME_MODE_SECONDS, float(position["seconds"])
-        else:
-            mode, value = LUMA_KEYFRAME_MODE_FRACTION, float(position["fraction"])
-        chain.add(LumaRay32KeyframeItem(image=image, mode=mode, value=value))
-        return IO.NodeOutput(chain)
-
-
-class LumaRay32KeyframesToVideoNode(IO.ComfyNode):
-    @classmethod
-    def define_schema(cls) -> IO.Schema:
-        return IO.Schema(
-            node_id="LumaRay32KeyframesToVideoNode",
-            display_name="Luma Ray 3.2 Keyframes to Video",
-            category="partner/video/Luma",
-            description="Generate a video that interpolates through a sequence of guide images, each anchored to a "
-            "position on the timeline, using Luma Ray 3.2. Build the sequence with Luma Ray 3.2 Keyframe nodes "
-            "(at least 2).",
-            inputs=[
-                IO.String.Input("prompt", multiline=True, default="", tooltip="Text prompt for the video generation."),
-                IO.Combo.Input("resolution", options=["360p", "540p", "720p", "1080p"], default="720p"),
-                IO.Combo.Input("duration", options=["5s", "10s"]),
-                _ray32_seed_input(),
-                IO.Custom(LumaIO.LUMA_RAY32_KEYFRAME).Input(
-                    "keyframes",
-                    tooltip="Keyframe sequence from Luma Ray 3.2 Keyframe nodes (at least 2).",
-                ),
-            ],
-            outputs=[IO.Video.Output(), IO.String.Output(display_name="generation_id")],
-            hidden=[
-                IO.Hidden.auth_token_comfy_org,
-                IO.Hidden.api_key_comfy_org,
-                IO.Hidden.unique_id,
-            ],
-            is_api_node=True,
-            price_badge=_BADGE_RAY32_VIDEO,
-        )
-
-    @classmethod
-    async def execute(
-        cls,
-        prompt: str,
-        resolution: str,
-        duration: str,
-        seed: int,
-        keyframes: LumaRay32KeyframeChain | None = None,
-    ) -> IO.NodeOutput:
-        validate_string(prompt, strip_whitespace=True, min_length=1, max_length=6000)
-        items = keyframes.items if keyframes is not None else []
-        if len(items) < 2:
-            raise ValueError(
-                "Connect at least 2 Luma Ray 3.2 Keyframe nodes "
-                "(use Luma Ray 3.2 Image to Video for a single start/end frame)."
-            )
-        if len(items) > 64:
-            raise ValueError(f"Ray 3.2 supports at most 64 keyframes; got {len(items)}.")
-        maxframe = 120 if duration == "5s" else 240
-        duration_seconds = maxframe / 24  # 5.0 or 10.0
-        # Resolve each keyframe to an output-frame index, then order by position
-        # (so the user can chain keyframes in any order — the position is what places them)
-        placed: list[tuple[int, torch.Tensor]] = []
-        for item in items:
-            if item.mode == LUMA_KEYFRAME_MODE_SECONDS:
-                if item.value > duration_seconds:
-                    raise ValueError(
-                        f"Keyframe position {item.value:g}s is past the end of the {duration} video; "
-                        f"use 0-{duration_seconds:g}s (or switch the keyframe to fraction mode)."
-                    )
-                idx = round(item.value * 24)
-            else:
-                idx = round(item.value * maxframe)
-            placed.append((max(0, min(maxframe, idx)), item.image))
-        placed.sort(key=lambda p: p[0])
-        indexes = [idx for idx, _ in placed]
-        for a, b in zip(indexes, indexes[1:]):
-            if a == b:
-                raise ValueError(
-                    f"Two keyframes resolve to the same output frame ({a}) for a {duration} video "
-                    f"(valid range 0-{maxframe}); give each keyframe a distinct position."
-                )
-        refs: list[Luma2ImageRef] = []
-        for _, image in placed:
-            url = await upload_image_to_comfyapi(cls, image, mime_type="image/png")
-            refs.append(Luma2ImageRef(url=url))
-        request = Luma2GenerationRequest(
-            prompt=prompt,
-            model="ray-3.2",
-            type="video",
-            video=Luma2VideoOptions(resolution=resolution, duration=duration, keyframes=refs, keyframe_indexes=indexes),
-        )
-        return await _ray32_generate(cls, request)
-
-
-class LumaRay32VideoEditNode(IO.ComfyNode):
-    @classmethod
-    def define_schema(cls) -> IO.Schema:
-        return IO.Schema(
-            node_id="LumaRay32VideoEditNode",
-            display_name="Luma Ray 3.2 Video Edit",
-            category="partner/video/Luma",
-            description="Re-render an existing video under a new prompt using Luma Ray 3.2 (restyle, relight, add "
-            "or remove elements) while keeping the original motion. Source video up to 18 seconds; the edited "
-            "video keeps the source's length.",
-            inputs=[
-                IO.Video.Input("video", tooltip="Source video to edit. Up to 18 seconds."),
-                IO.String.Input("prompt", multiline=True, default="", tooltip="Describes the desired edit."),
-                IO.Combo.Input("resolution", options=["360p", "540p", "720p", "1080p"], default="720p"),
-                IO.Combo.Input(
-                    "strength",
-                    options=[
-                        "auto",
-                        "adhere_1",
-                        "adhere_2",
-                        "adhere_3",
-                        "flex_1",
-                        "flex_2",
-                        "flex_3",
-                        "reimagine_1",
-                        "reimagine_2",
-                        "reimagine_3",
-                    ],
-                    default="auto",
-                    tooltip="How strongly to preserve vs. reimagine the source. 'auto' lets Ray 3.2 choose; "
-                    "adhere_* preserves the most, flex_* is balanced, reimagine_* changes the most.",
-                ),
-                _ray32_seed_input(),
-            ],
-            outputs=[
-                IO.Video.Output(),
-                IO.String.Output(display_name="generation_id"),
-            ],
-            hidden=[
-                IO.Hidden.auth_token_comfy_org,
-                IO.Hidden.api_key_comfy_org,
-                IO.Hidden.unique_id,
-            ],
-            is_api_node=True,
-            price_badge=_BADGE_RAY32_EDIT,
-        )
-
-    @classmethod
-    async def execute(
-        cls, video: Input.Video, prompt: str, resolution: str, strength: str, seed: int
-    ) -> IO.NodeOutput:
-        validate_string(prompt, strip_whitespace=True, min_length=1, max_length=6000)
-        try:
-            duration = "5s" if video.get_duration() <= 5.0 else "10s"
-        except Exception:
-            duration = "10s"
-        source_url = await upload_video_to_comfyapi(cls, video, max_duration=18)
-        edit = Luma2VideoEdit(auto_controls=True) if strength == "auto" else Luma2VideoEdit(strength=strength)
-        request = Luma2GenerationRequest(
-            prompt=prompt,
-            model="ray-3.2",
-            type="video_edit",
-            source=Luma2ImageRef(url=source_url, media_type="video/mp4"),
-            video=Luma2VideoOptions(resolution=resolution, duration=duration, edit=edit),
-        )
-        return await _ray32_generate(cls, request)
-
-
-class LumaRay32VideoReframeNode(IO.ComfyNode):
-    @classmethod
-    def define_schema(cls) -> IO.Schema:
-        return IO.Schema(
-            node_id="LumaRay32VideoReframeNode",
-            display_name="Luma Ray 3.2 Video Reframe",
-            category="partner/video/Luma",
-            description="Change the aspect ratio of an existing video, using Luma Ray 3.2 to fill the newly "
-            "exposed canvas areas. Source video up to 30 seconds. Billed per second of output.",
-            inputs=[
-                IO.Video.Input("video", tooltip="Source video to reframe. Up to 30 seconds."),
-                IO.String.Input(
-                    "prompt",
-                    multiline=True,
-                    default="",
-                    tooltip="Describes how the newly exposed canvas areas should be filled.",
-                ),
-                IO.Combo.Input("aspect_ratio", options=["16:9", "9:16", "1:1", "4:3", "3:4", "21:9"]),
-                IO.Combo.Input("resolution", options=["360p", "540p", "720p", "1080p"], default="720p"),
-                _ray32_seed_input(),
-            ],
-            outputs=[
-                IO.Video.Output(),
-                IO.String.Output(display_name="generation_id"),
-            ],
-            hidden=[
-                IO.Hidden.auth_token_comfy_org,
-                IO.Hidden.api_key_comfy_org,
-                IO.Hidden.unique_id,
-            ],
-            is_api_node=True,
-            price_badge=_BADGE_RAY32_REFRAME,
-        )
-
-    @classmethod
-    async def execute(
-        cls, video: Input.Video, prompt: str, aspect_ratio: str, resolution: str, seed: int
-    ) -> IO.NodeOutput:
-        validate_string(prompt, strip_whitespace=False, min_length=1, max_length=6000)
-        if resolution == "1080p" and aspect_ratio in {"9:16", "3:4"}:
-            raise ValueError("1080p is not available for vertical aspect ratios (9:16, 3:4) when reframing.")
-        source_url = await upload_video_to_comfyapi(cls, video, max_duration=30)
-        request = Luma2GenerationRequest(
-            prompt=prompt,
-            model="ray-3.2",
-            type="video_reframe",
-            aspect_ratio=aspect_ratio,
-            source=Luma2ImageRef(url=source_url, media_type="video/mp4"),
-            video=Luma2VideoOptions(resolution=resolution),
-        )
-        return await _ray32_generate(cls, request)
-
-
-class LumaRay32ExtendVideoNode(IO.ComfyNode):
-    @classmethod
-    def define_schema(cls) -> IO.Schema:
-        return IO.Schema(
-            node_id="LumaRay32ExtendVideoNode",
-            display_name="Luma Ray 3.2 Extend Video",
-            category="partner/video/Luma",
-            description="Extend a previous Ray 3.2 generation forward (continue after it) or backward (lead-in "
-            "before it). Connect the generation_id output of a prior Luma Ray 3.2 node."
-            " Extensions are always 5 seconds.",
-            inputs=[
-                IO.String.Input(
-                    "source_generation_id",
-                    default="",
-                    tooltip="generation_id of the prior Ray 3.2 video to extend."
-                    " Connect the generation_id output of another Luma Ray 3.2 node.",
-                ),
-                IO.DynamicCombo.Input(
-                    "direction",
-                    options=[
-                        IO.DynamicCombo.Option(
-                            "Forward (continue after)",
-                            [
-                                IO.Boolean.Input(
-                                    "loop",
-                                    default=False,
-                                    tooltip="Loop the extended video seamlessly (forward extend only).",
-                                ),
-                            ],
-                        ),
-                        IO.DynamicCombo.Option("Backward (lead-in before)", []),
-                    ],
-                    tooltip="Forward continues after the prior clip; backward is prepended before it.",
-                ),
-                IO.String.Input("prompt", multiline=True, default="", tooltip="Text prompt for the new content."),
-                IO.Combo.Input("resolution", options=["540p", "720p", "1080p"], default="720p"),
-                _ray32_seed_input(),
-            ],
-            outputs=[
-                IO.Video.Output(),
-                IO.String.Output(display_name="generation_id"),
-            ],
-            hidden=[
-                IO.Hidden.auth_token_comfy_org,
-                IO.Hidden.api_key_comfy_org,
-                IO.Hidden.unique_id,
-            ],
-            is_api_node=True,
-            price_badge=_BADGE_RAY32_VIDEO_5S,
-        )
-
-    @classmethod
-    async def execute(
-        cls, source_generation_id: str, direction: dict, prompt: str, resolution: str, seed: int
-    ) -> IO.NodeOutput:
-        validate_string(prompt, strip_whitespace=False, min_length=1, max_length=6000)
-        gen_id = (source_generation_id or "").strip()
-        if not gen_id:
-            raise ValueError(
-                "source_generation_id is required (connect the generation_id output of a prior Luma Ray 3.2 node)."
-            )
-        video = Luma2VideoOptions(resolution=resolution, duration="5s")
-        ref = Luma2ImageRef(generation_id=gen_id)
-        if direction["direction"] == "Forward (continue after)":
-            video.start_frame = ref
-            if direction.get("loop"):
-                video.loop = True
-        else:
-            video.end_frame = ref
-        request = Luma2GenerationRequest(prompt=prompt, model="ray-3.2", type="video", video=video)
-        return await _ray32_generate(cls, request)
+        return IO.NodeOutput(await _luma2_submit_and_poll(cls, request))
 
 
 class LumaExtension(ComfyExtension):
@@ -1481,13 +944,6 @@ class LumaExtension(ComfyExtension):
             LumaConceptsNode,
             LumaImageNode,
             LumaImageEditNode,
-            LumaRay32TextToVideoNode,
-            LumaRay32ImageToVideoNode,
-            LumaRay32KeyframeNode,
-            LumaRay32KeyframesToVideoNode,
-            LumaRay32VideoEditNode,
-            LumaRay32VideoReframeNode,
-            LumaRay32ExtendVideoNode,
         ]
 
 

comfy_extras/nodes_dataset.py

@@ -1583,7 +1583,7 @@ class LoadTrainingDataset(io.ComfyNode):
             shard_path = os.path.join(dataset_dir, shard_file)
 
             with open(shard_path, "rb") as f:
-                shard_data = torch.load(f)
+                shard_data = torch.load(f, weights_only=True)
 
             all_latents.extend(shard_data["latents"])
             all_conditioning.extend(shard_data["conditioning"])