Merge 6e434c785b into 6cc814437f

comfy/ldm/ideogram4/model.py

@@ -226,27 +226,85 @@ class Ideogram4Transformer2DModel(Ideogram4Transformer):
         x = x.permute(0, 5, 3, 4, 1, 2)  # (B, c, pi, pj, gh, gw)
         return x.reshape(B, C, gh, gw)
 
-    def _image_position_ids(self, gh, gw, device):
-        h_idx = torch.arange(gh, device=device).view(-1, 1).expand(gh, gw).reshape(-1)
-        w_idx = torch.arange(gw, device=device).view(1, -1).expand(gh, gw).reshape(-1)
-        t_idx = torch.zeros_like(h_idx)
+    def _image_position_ids(self, gh, gw, device, index=0, h_offset=0, w_offset=0):
+        h_idx = torch.arange(gh, device=device).view(-1, 1).expand(gh, gw).reshape(-1) + h_offset
+        w_idx = torch.arange(gw, device=device).view(1, -1).expand(gh, gw).reshape(-1) + w_offset
+        t_idx = torch.full_like(h_idx, index)
         return torch.stack([t_idx, h_idx, w_idx], dim=1) + IMAGE_POSITION_OFFSET  # (L_img, 3)
 
-    def _run_conditional(self, x_chunk, context_chunk, attn_mask_chunk, t_chunk, gh, gw, transformer_options):
+    def _run_conditional(self, x_chunk, context_chunk, attn_mask_chunk, t_chunk, gh, gw, transformer_options, ref_latents=None, ref_latents_method="index"):
         B = x_chunk.shape[0]
         device = x_chunk.device
         img_tokens = self._img_to_tokens(x_chunk)
         L_img = img_tokens.shape[1]
         L_text = context_chunk.shape[1]
-        L = L_text + L_img
         latent_dim = img_tokens.shape[-1]
 
+        ref_tokens_list = []
+        ref_pos_ids_list = []
+        ref_num_tokens = []
+
+        if ref_latents is not None:
+            h = 0
+            w = 0
+            index = 0
+            index_ref_method = (ref_latents_method == "index") or (ref_latents_method == "index_timestep_zero")
+            negative_ref_method = ref_latents_method == "negative_index"
+
+            for ref in ref_latents:
+                ref_b, ref_c, ref_h, ref_w = ref.shape
+                ref_gh = ref_h
+                ref_gw = ref_w
+
+                if index_ref_method:
+                    index += 1
+                    gh_offset = 0
+                    gw_offset = 0
+                elif negative_ref_method:
+                    index -= 1
+                    gh_offset = 0
+                    gw_offset = 0
+                else: # offset/default
+                    index = 1
+                    gh_offset = 0
+                    gw_offset = 0
+                    if ref_gh + h > ref_gw + w:
+                        gw_offset = w
+                    else:
+                        gh_offset = h
+                    h = max(h, ref_gh + gh_offset)
+                    w = max(w, ref_gw + gw_offset)
+
+                ref_tokens = self._img_to_tokens(ref)
+                ref_tokens_list.append(ref_tokens)
+                ref_num_tokens.append(ref_tokens.shape[1])
+
+                ref_pos = self._image_position_ids(ref_gh, ref_gw, device, index=index, h_offset=gh_offset, w_offset=gw_offset)
+                ref_pos_ids_list.append(ref_pos)
+
+            transformer_options = transformer_options.copy()
+            transformer_options["reference_image_num_tokens"] = ref_num_tokens
+
+        L_ref = sum(t.shape[1] for t in ref_tokens_list) if ref_tokens_list else 0
+        L = L_text + L_img + L_ref
+
         x_full = torch.zeros(B, L, latent_dim, dtype=img_tokens.dtype, device=device)
-        x_full[:, L_text:] = img_tokens
+        x_full[:, L_text:L_text+L_img] = img_tokens
+
+        curr_idx = L_text + L_img
+        for ref_tokens in ref_tokens_list:
+            ref_len = ref_tokens.shape[1]
+            x_full[:, curr_idx:curr_idx+ref_len] = ref_tokens
+            curr_idx += ref_len
 
         text_pos = torch.arange(L_text, device=device).view(-1, 1).expand(L_text, 3)
         img_pos = self._image_position_ids(gh, gw, device)
-        position_ids = torch.cat([text_pos, img_pos], dim=0).unsqueeze(0).expand(B, L, 3)
+
+        pos_ids_all = [text_pos, img_pos]
+        for ref_pos in ref_pos_ids_list:
+            pos_ids_all.append(ref_pos)
+
+        position_ids = torch.cat(pos_ids_all, dim=0).unsqueeze(0).expand(B, L, 3)
 
         indicator = torch.empty(B, L, dtype=torch.long, device=device)
         indicator[:, :L_text] = LLM_TOKEN_INDICATOR
@@ -263,20 +321,84 @@ class Ideogram4Transformer2DModel(Ideogram4Transformer):
 
         out = self._backbone(context_chunk, x_full, t_chunk, position_ids, attn_mask, indicator,
                              transformer_options=transformer_options)
-        return self._tokens_to_img(out[:, L_text:], gh, gw)
+        return self._tokens_to_img(out[:, L_text:L_text+L_img], gh, gw)
 
-    def _run_image_only(self, x_chunk, t_chunk, gh, gw, transformer_options):
+    def _run_image_only(self, x_chunk, t_chunk, gh, gw, transformer_options, ref_latents=None, ref_latents_method="index"):
         B = x_chunk.shape[0]
         device = x_chunk.device
         img_tokens = self._img_to_tokens(x_chunk)
         L_img = img_tokens.shape[1]
+        latent_dim = img_tokens.shape[-1]
 
-        position_ids = self._image_position_ids(gh, gw, device).unsqueeze(0).expand(B, L_img, 3)
-        indicator = torch.full((B, L_img), OUTPUT_IMAGE_INDICATOR, dtype=torch.long, device=device)
+        ref_tokens_list = []
+        ref_pos_ids_list = []
+        ref_num_tokens = []
+
+        if ref_latents is not None:
+            h = 0
+            w = 0
+            index = 0
+            index_ref_method = (ref_latents_method == "index") or (ref_latents_method == "index_timestep_zero")
+            negative_ref_method = ref_latents_method == "negative_index"
+
+            for ref in ref_latents:
+                ref_b, ref_c, ref_h, ref_w = ref.shape
+                ref_gh = ref_h
+                ref_gw = ref_w
+
+                if index_ref_method:
+                    index += 1
+                    gh_offset = 0
+                    gw_offset = 0
+                elif negative_ref_method:
+                    index -= 1
+                    gh_offset = 0
+                    gw_offset = 0
+                else: # offset/default
+                    index = 1
+                    gh_offset = 0
+                    gw_offset = 0
+                    if ref_gh + h > ref_gw + w:
+                        gw_offset = w
+                    else:
+                        gh_offset = h
+                    h = max(h, ref_gh + gh_offset)
+                    w = max(w, ref_gw + gw_offset)
+
+                ref_tokens = self._img_to_tokens(ref)
+                ref_tokens_list.append(ref_tokens)
+                ref_num_tokens.append(ref_tokens.shape[1])
+
+                ref_pos = self._image_position_ids(ref_gh, ref_gw, device, index=index, h_offset=gh_offset, w_offset=gw_offset)
+                ref_pos_ids_list.append(ref_pos)
+
+            transformer_options = transformer_options.copy()
+            transformer_options["reference_image_num_tokens"] = ref_num_tokens
+
+        L_ref = sum(t.shape[1] for t in ref_tokens_list) if ref_tokens_list else 0
+        L_img_total = L_img + L_ref
+
+        x_full = torch.zeros(B, L_img_total, latent_dim, dtype=img_tokens.dtype, device=device)
+        x_full[:, :L_img] = img_tokens
+
+        curr_idx = L_img
+        for ref_tokens in ref_tokens_list:
+            ref_len = ref_tokens.shape[1]
+            x_full[:, curr_idx:curr_idx+ref_len] = ref_tokens
+            curr_idx += ref_len
+
+        img_pos = self._image_position_ids(gh, gw, device)
+
+        pos_ids_all = [img_pos]
+        for ref_pos in ref_pos_ids_list:
+            pos_ids_all.append(ref_pos)
+
+        position_ids = torch.cat(pos_ids_all, dim=0).unsqueeze(0).expand(B, L_img_total, 3)
+        indicator = torch.full((B, L_img_total), OUTPUT_IMAGE_INDICATOR, dtype=torch.long, device=device)
 
         # Image-only sequence is a single segment -> no mask, full attention, no LLM context.
-        out = self._backbone(None, img_tokens, t_chunk, position_ids, None, indicator, transformer_options=transformer_options)
-        return self._tokens_to_img(out, gh, gw)
+        out = self._backbone(None, x_full, t_chunk, position_ids, None, indicator, transformer_options=transformer_options)
+        return self._tokens_to_img(out[:, :L_img], gh, gw)
 
     def forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, **kwargs):
         return comfy.patcher_extension.WrapperExecutor.new_class_executor(
@@ -290,8 +412,11 @@ class Ideogram4Transformer2DModel(Ideogram4Transformer):
 
         timesteps = 1.0 - timesteps
 
+        ref_latents = kwargs.get("ref_latents", None)
+        ref_latents_method = kwargs.get("ref_latents_method", "index")
+
         # unconditional pass
         if context is None:
-            return -self._run_image_only(x, timesteps, gh, gw, transformer_options)
+            return -self._run_image_only(x, timesteps, gh, gw, transformer_options, ref_latents=ref_latents, ref_latents_method=ref_latents_method)
 
-        return -self._run_conditional(x, context, attention_mask, timesteps, gh, gw, transformer_options)
+        return -self._run_conditional(x, context, attention_mask, timesteps, gh, gw, transformer_options, ref_latents=ref_latents, ref_latents_method=ref_latents_method)

comfy/model_base.py

@@ -2267,6 +2267,7 @@ class QwenImage(BaseModel):
 class Ideogram4(BaseModel):
     def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.ideogram4.model.Ideogram4Transformer2DModel)
+        self.memory_usage_factor_conds = ("ref_latents",)
 
     def extra_conds(self, **kwargs):
         out = super().extra_conds(**kwargs)
@@ -2277,6 +2278,24 @@ class Ideogram4(BaseModel):
         cross_attn = kwargs.get("cross_attn", None)
         if cross_attn is not None:
             out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
+
+        ref_latents = kwargs.get("reference_latents", None)
+        if ref_latents is not None:
+            latents = []
+            for lat in ref_latents:
+                latents.append(self.process_latent_in(lat))
+            out['ref_latents'] = comfy.conds.CONDList(latents)
+
+            ref_latents_method = kwargs.get("reference_latents_method", None)
+            if ref_latents_method is not None:
+                out['ref_latents_method'] = comfy.conds.CONDConstant(ref_latents_method)
+        return out
+
+    def extra_conds_shapes(self, **kwargs):
+        out = {}
+        ref_latents = kwargs.get("reference_latents", None)
+        if ref_latents is not None:
+            out['ref_latents'] = list([1, 128, sum(map(lambda a: math.prod(a.size()[2:]), ref_latents))])
         return out
 
 class Krea2(BaseModel):

comfy/sd.py

@@ -1255,7 +1255,10 @@ class VAE:
             return None
 
     def is_dynamic(self):
-        return self.patcher.is_dynamic()
+        # A VAE built from a state dict with no detectable VAE weights returns early
+        # from __init__ ("No VAE weights detected") before self.patcher is assigned.
+        patcher = getattr(self, "patcher", None)
+        return patcher is not None and patcher.is_dynamic()
 
 class StyleModel:
     def __init__(self, model, device="cpu"):

comfy_api_nodes/apis/bytedance.py

@@ -24,8 +24,8 @@ class Seedream4TaskCreationRequest(BaseModel):
     image: list[str] | None = Field(None, description="Image URLs")
     size: str = Field(...)
     seed: int = Field(..., ge=0, le=2147483647)
-    sequential_image_generation: str = Field("disabled")
-    sequential_image_generation_options: Seedream4Options = Field(Seedream4Options(max_images=15))
+    sequential_image_generation: str | None = Field("disabled")
+    sequential_image_generation_options: Seedream4Options | None = Field(Seedream4Options(max_images=15))
     watermark: bool = Field(False)
     output_format: str | None = None
 
@@ -261,6 +261,19 @@ _PRESETS_SEEDREAM_4K = [
 
 _CUSTOM_PRESET = [("Custom", None, None)]
 
+_PRESETS_SEEDREAM_2K_PRO = [
+    ("(2K) 2048x2048 (1:1)", 2048, 2048),
+    ("(2K) 1728x2304 (3:4)", 1728, 2304),
+    ("(2K) 2304x1728 (4:3)", 2304, 1728),
+    # ("(2K) 2848x1600 (16:9)", 2848, 1600),  # 4,556,800 px - temporarily unavailable
+    # ("(2K) 1600x2848 (9:16)", 1600, 2848),  # 4,556,800 px - temporarily unavailable
+    ("(2K) 1664x2496 (2:3)", 1664, 2496),
+    ("(2K) 2496x1664 (3:2)", 2496, 1664),
+    # ("(2K) 3136x1344 (21:9)", 3136, 1344),  # 4,214,784 px - temporarily unavailable
+]
+RECOMMENDED_PRESETS_SEEDREAM_5_PRO = (
+    _PRESETS_SEEDREAM_1K + _PRESETS_SEEDREAM_2K_PRO + _CUSTOM_PRESET
+)
 RECOMMENDED_PRESETS_SEEDREAM_5_LITE = (
     _PRESETS_SEEDREAM_2K + _PRESETS_SEEDREAM_3K + _PRESETS_SEEDREAM_4K + _CUSTOM_PRESET
 )

comfy_api_nodes/nodes_bytedance.py

@@ -16,6 +16,7 @@ from comfy_api_nodes.apis.bytedance import (
     RECOMMENDED_PRESETS_SEEDREAM_4_0,
     RECOMMENDED_PRESETS_SEEDREAM_4_5,
     RECOMMENDED_PRESETS_SEEDREAM_5_LITE,
+    RECOMMENDED_PRESETS_SEEDREAM_5_PRO,
     SEEDANCE2_REF_VIDEO_PIXEL_LIMITS,
     VIDEO_TASKS_EXECUTION_TIME,
     GetAssetResponse,
@@ -80,12 +81,14 @@ _VERIFICATION_POLL_TIMEOUT_SEC = 120
 _VERIFICATION_POLL_INTERVAL_SEC = 3
 
 SEEDREAM_MODELS = {
+    "seedream 5.0 pro": "seedream-5-0-pro-260628",
     "seedream 5.0 lite": "seedream-5-0-260128",
     "seedream-4-5-251128": "seedream-4-5-251128",
     "seedream-4-0-250828": "seedream-4-0-250828",
 }
 
 SEEDREAM_PRESETS = {
+    "seedream-5-0-pro-260628": RECOMMENDED_PRESETS_SEEDREAM_5_PRO,
     "seedream-5-0-260128": RECOMMENDED_PRESETS_SEEDREAM_5_LITE,
     "seedream-4-5-251128": RECOMMENDED_PRESETS_SEEDREAM_4_5,
     "seedream-4-0-250828": RECOMMENDED_PRESETS_SEEDREAM_4_0,
@@ -743,8 +746,15 @@ class ByteDanceSeedreamNode(IO.ComfyNode):
         return IO.NodeOutput(torch.cat([await download_url_to_image_tensor(i) for i in urls]))
 
 
-def _seedream_model_inputs(*, max_ref_images: int, presets: list):
-    return [
+def _seedream_model_inputs(
+    *,
+    max_ref_images: int,
+    presets: list,
+    max_width: int = 6240,
+    max_height: int = 4992,
+    supports_batch: bool = True,
+):
+    inputs = [
         IO.Combo.Input(
             "size_preset",
             options=[label for label, _, _ in presets],
@@ -754,7 +764,7 @@ def _seedream_model_inputs(*, max_ref_images: int, presets: list):
             "width",
             default=2048,
             min=1024,
-            max=6240,
+            max=max_width,
             step=2,
             tooltip="Custom width for image. Value is working only if `size_preset` is set to `Custom`",
         ),
@@ -762,22 +772,27 @@ def _seedream_model_inputs(*, max_ref_images: int, presets: list):
             "height",
             default=2048,
             min=1024,
-            max=4992,
+            max=max_height,
             step=2,
             tooltip="Custom height for image. Value is working only if `size_preset` is set to `Custom`",
         ),
-        IO.Int.Input(
-            "max_images",
-            default=1,
-            min=1,
-            max=max_ref_images,
-            step=1,
-            display_mode=IO.NumberDisplay.number,
-            tooltip="Maximum number of images to generate. With 1, exactly one image is produced. "
-            "With >1, the model generates between 1 and max_images related images "
-            "(e.g., story scenes, character variations). "
-            "Total images (input + generated) cannot exceed 15.",
-        ),
+    ]
+    if supports_batch:
+        inputs.append(
+            IO.Int.Input(
+                "max_images",
+                default=1,
+                min=1,
+                max=max_ref_images,
+                step=1,
+                display_mode=IO.NumberDisplay.number,
+                tooltip="Maximum number of images to generate. With 1, exactly one image is produced. "
+                "With >1, the model generates between 1 and max_images related images "
+                "(e.g., story scenes, character variations). "
+                "Total images (input + generated) cannot exceed 15.",
+            )
+        )
+    inputs.append(
         IO.Autogrow.Input(
             "images",
             template=IO.Autogrow.TemplateNames(
@@ -787,14 +802,18 @@ def _seedream_model_inputs(*, max_ref_images: int, presets: list):
             ),
             tooltip=f"Optional reference image(s) for image-to-image or multi-reference generation. "
             f"Up to {max_ref_images} images.",
-        ),
-        IO.Boolean.Input(
-            "fail_on_partial",
-            default=False,
-            tooltip="If enabled, abort execution if any requested images are missing or return an error.",
-            advanced=True,
-        ),
-    ]
+        )
+    )
+    if supports_batch:
+        inputs.append(
+            IO.Boolean.Input(
+                "fail_on_partial",
+                default=False,
+                tooltip="If enabled, abort execution if any requested images are missing or return an error.",
+                advanced=True,
+            )
+        )
+    return inputs
 
 
 class ByteDanceSeedreamNodeV2(IO.ComfyNode):
@@ -816,6 +835,16 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
                 IO.DynamicCombo.Input(
                     "model",
                     options=[
+                        IO.DynamicCombo.Option(
+                            "seedream 5.0 pro",
+                            _seedream_model_inputs(
+                                max_ref_images=10,
+                                presets=RECOMMENDED_PRESETS_SEEDREAM_5_PRO,
+                                max_width=3136,
+                                max_height=2496,
+                                supports_batch=False,
+                            ),
+                        ),
                         IO.DynamicCombo.Option(
                             "seedream 5.0 lite",
                             _seedream_model_inputs(max_ref_images=14, presets=RECOMMENDED_PRESETS_SEEDREAM_5_LITE),
@@ -857,15 +886,27 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
             ],
             is_api_node=True,
             price_badge=IO.PriceBadge(
-                depends_on=IO.PriceBadgeDepends(widgets=["model"]),
+                depends_on=IO.PriceBadgeDepends(
+                    widgets=["model", "model.size_preset", "model.width", "model.height"]
+                ),
                 expr="""
                 (
-                  $price := $contains(widgets.model, "5.0 lite") ? 0.035 :
-                            $contains(widgets.model, "4-5") ? 0.04 : 0.03;
+                  $sp := $lookup(widgets, "model.size_preset");
+                  $px := $lookup(widgets, "model.width") * $lookup(widgets, "model.height");
+                  $isPro := $contains(widgets.model, "5.0 pro");
+                  $price := $isPro
+                              ? (
+                                  $contains($sp, "custom")
+                                    ? ($px <= 2360000 ? 0.045 : 0.09)
+                                    : ($contains($sp, "1k") ? 0.045 : 0.09)
+                                )
+                              : $contains(widgets.model, "5.0 lite") ? 0.035
+                              : $contains(widgets.model, "4-5") ? 0.04
+                              : 0.03;
                   {
-                    "type":"usd",
+                    "type": "usd",
                     "usd": $price,
-                    "format": { "suffix":" x images/Run", "approximate": true }
+                    "format": { "suffix": $isPro ? "/Image" : " x images/Run", "approximate": true }
                   }
                 )
                 """,
@@ -883,6 +924,7 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
         validate_string(prompt, strip_whitespace=True, min_length=1)
         model_id = SEEDREAM_MODELS[model["model"]]
         presets = SEEDREAM_PRESETS[model_id]
+        is_pro = "seedream-5-0-pro" in model_id
 
         size_preset = model.get("size_preset", presets[0][0])
         width = model.get("width", 2048)
@@ -902,19 +944,29 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
 
         out_num_pixels = w * h
         mp_provided = out_num_pixels / 1_000_000.0
-        if ("seedream-4-5" in model_id or "seedream-5-0" in model_id) and out_num_pixels < 3686400:
-            raise ValueError(
-                f"Minimum image resolution for the selected model is 3.68MP, but {mp_provided:.2f}MP provided."
-            )
-        if "seedream-4-0" in model_id and out_num_pixels < 921600:
-            raise ValueError(
-                f"Minimum image resolution that the selected model can generate is 0.92MP, "
-                f"but {mp_provided:.2f}MP provided."
-            )
-        if out_num_pixels > 16_777_216:
-            raise ValueError(
-                f"Maximum image resolution for the selected model is 16.78MP, but {mp_provided:.2f}MP provided."
-            )
+        if is_pro:
+            if out_num_pixels < 921_600:
+                raise ValueError(
+                    f"Minimum image resolution for the selected model is 0.92MP, but {mp_provided:.2f}MP provided."
+                )
+            if out_num_pixels > 4_194_304:
+                raise ValueError(
+                    f"Maximum image resolution for the selected model is 4.19MP, but {mp_provided:.2f}MP provided."
+                )
+        else:
+            if ("seedream-4-5" in model_id or "seedream-5-0" in model_id) and out_num_pixels < 3_686_400:
+                raise ValueError(
+                    f"Minimum image resolution for the selected model is 3.68MP, but {mp_provided:.2f}MP provided."
+                )
+            if "seedream-4-0" in model_id and out_num_pixels < 921_600:
+                raise ValueError(
+                    f"Minimum image resolution that the selected model can generate is 0.92MP, "
+                    f"but {mp_provided:.2f}MP provided."
+                )
+            if out_num_pixels > 16_777_216:
+                raise ValueError(
+                    f"Maximum image resolution for the selected model is 16.78MP, but {mp_provided:.2f}MP provided."
+                )
 
         image_tensors: list[Input.Image] = [t for t in images_dict.values() if t is not None]
         n_input_images = sum(get_number_of_images(t) for t in image_tensors)
@@ -950,8 +1002,8 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
                 image=reference_images_urls,
                 size=f"{w}x{h}",
                 seed=seed,
-                sequential_image_generation=sequential_image_generation,
-                sequential_image_generation_options=Seedream4Options(max_images=max_images),
+                sequential_image_generation=None if is_pro else sequential_image_generation,
+                sequential_image_generation_options=None if is_pro else Seedream4Options(max_images=max_images),
                 watermark=watermark,
             ),
         )

requirements.txt

@@ -1,5 +1,5 @@
 comfyui-frontend-package==1.45.20
-comfyui-workflow-templates==0.11.2
+comfyui-workflow-templates==0.11.6
 comfyui-embedded-docs==0.5.7
 torch
 torchsde