Merge branch 'comfyanonymous:master' into chroma-support

comfy/comfy_types/node_typing.py

@@ -99,59 +99,59 @@ class InputTypeOptions(TypedDict):
     Comfy Docs: https://docs.comfy.org/custom-nodes/backend/datatypes
     """
 
-    default: bool | str | float | int | list | tuple
+    default: NotRequired[bool | str | float | int | list | tuple]
     """The default value of the widget"""
-    defaultInput: bool
+    defaultInput: NotRequired[bool]
     """@deprecated in v1.16 frontend. v1.16 frontend allows input socket and widget to co-exist.
     - defaultInput on required inputs should be dropped.
     - defaultInput on optional inputs should be replaced with forceInput.
     Ref: https://github.com/Comfy-Org/ComfyUI_frontend/pull/3364
     """
-    forceInput: bool
+    forceInput: NotRequired[bool]
     """Forces the input to be an input slot rather than a widget even a widget is available for the input type."""
-    lazy: bool
+    lazy: NotRequired[bool]
     """Declares that this input uses lazy evaluation"""
-    rawLink: bool
+    rawLink: NotRequired[bool]
     """When a link exists, rather than receiving the evaluated value, you will receive the link (i.e. `["nodeId", <outputIndex>]`). Designed for node expansion."""
-    tooltip: str
+    tooltip: NotRequired[str]
     """Tooltip for the input (or widget), shown on pointer hover"""
     # class InputTypeNumber(InputTypeOptions):
     # default: float | int
-    min: float
+    min: NotRequired[float]
     """The minimum value of a number (``FLOAT`` | ``INT``)"""
-    max: float
+    max: NotRequired[float]
     """The maximum value of a number (``FLOAT`` | ``INT``)"""
-    step: float
+    step: NotRequired[float]
     """The amount to increment or decrement a widget by when stepping up/down (``FLOAT`` | ``INT``)"""
-    round: float
+    round: NotRequired[float]
     """Floats are rounded by this value (``FLOAT``)"""
     # class InputTypeBoolean(InputTypeOptions):
     # default: bool
-    label_on: str
+    label_on: NotRequired[str]
     """The label to use in the UI when the bool is True (``BOOLEAN``)"""
-    label_off: str
+    label_off: NotRequired[str]
     """The label to use in the UI when the bool is False (``BOOLEAN``)"""
     # class InputTypeString(InputTypeOptions):
     # default: str
-    multiline: bool
+    multiline: NotRequired[bool]
     """Use a multiline text box (``STRING``)"""
-    placeholder: str
+    placeholder: NotRequired[str]
     """Placeholder text to display in the UI when empty (``STRING``)"""
     # Deprecated:
     # defaultVal: str
-    dynamicPrompts: bool
+    dynamicPrompts: NotRequired[bool]
     """Causes the front-end to evaluate dynamic prompts (``STRING``)"""
     # class InputTypeCombo(InputTypeOptions):
-    image_upload: bool
+    image_upload: NotRequired[bool]
     """Specifies whether the input should have an image upload button and image preview attached to it. Requires that the input's name is `image`."""
-    image_folder: Literal["input", "output", "temp"]
+    image_folder: NotRequired[Literal["input", "output", "temp"]]
     """Specifies which folder to get preview images from if the input has the ``image_upload`` flag.
     """
-    remote: RemoteInputOptions
+    remote: NotRequired[RemoteInputOptions]
     """Specifies the configuration for a remote input.
     Available after ComfyUI frontend v1.9.7
     https://github.com/Comfy-Org/ComfyUI_frontend/pull/2422"""
-    control_after_generate: bool
+    control_after_generate: NotRequired[bool]
     """Specifies whether a control widget should be added to the input, adding options to automatically change the value after each prompt is queued. Currently only used for INT and COMBO types."""
     options: NotRequired[list[str | int | float]]
     """COMBO type only. Specifies the selectable options for the combo widget.
@@ -169,15 +169,15 @@ class InputTypeOptions(TypedDict):
 class HiddenInputTypeDict(TypedDict):
     """Provides type hinting for the hidden entry of node INPUT_TYPES."""
 
-    node_id: Literal["UNIQUE_ID"]
+    node_id: NotRequired[Literal["UNIQUE_ID"]]
     """UNIQUE_ID is the unique identifier of the node, and matches the id property of the node on the client side. It is commonly used in client-server communications (see messages)."""
-    unique_id: Literal["UNIQUE_ID"]
+    unique_id: NotRequired[Literal["UNIQUE_ID"]]
     """UNIQUE_ID is the unique identifier of the node, and matches the id property of the node on the client side. It is commonly used in client-server communications (see messages)."""
-    prompt: Literal["PROMPT"]
+    prompt: NotRequired[Literal["PROMPT"]]
     """PROMPT is the complete prompt sent by the client to the server. See the prompt object for a full description."""
-    extra_pnginfo: Literal["EXTRA_PNGINFO"]
+    extra_pnginfo: NotRequired[Literal["EXTRA_PNGINFO"]]
     """EXTRA_PNGINFO is a dictionary that will be copied into the metadata of any .png files saved. Custom nodes can store additional information in this dictionary for saving (or as a way to communicate with a downstream node)."""
-    dynprompt: Literal["DYNPROMPT"]
+    dynprompt: NotRequired[Literal["DYNPROMPT"]]
     """DYNPROMPT is an instance of comfy_execution.graph.DynamicPrompt. It differs from PROMPT in that it may mutate during the course of execution in response to Node Expansion."""
 
 
@@ -187,11 +187,11 @@ class InputTypeDict(TypedDict):
     Comfy Docs: https://docs.comfy.org/custom-nodes/backend/more_on_inputs
     """
 
-    required: dict[str, tuple[IO, InputTypeOptions]]
+    required: NotRequired[dict[str, tuple[IO, InputTypeOptions]]]
     """Describes all inputs that must be connected for the node to execute."""
-    optional: dict[str, tuple[IO, InputTypeOptions]]
+    optional: NotRequired[dict[str, tuple[IO, InputTypeOptions]]]
     """Describes inputs which do not need to be connected."""
-    hidden: HiddenInputTypeDict
+    hidden: NotRequired[HiddenInputTypeDict]
     """Offers advanced functionality and server-client communication.
 
     Comfy Docs: https://docs.comfy.org/custom-nodes/backend/more_on_inputs#hidden-inputs

comfy/ldm/hidream/model.py

@@ -8,25 +8,12 @@ from einops import repeat
 from comfy.ldm.lightricks.model import TimestepEmbedding, Timesteps
 import torch.nn.functional as F
 
-from comfy.ldm.flux.math import apply_rope
+from comfy.ldm.flux.math import apply_rope, rope
+from comfy.ldm.flux.layers import LastLayer
+
 from comfy.ldm.modules.attention import optimized_attention
 import comfy.model_management
-
+import comfy.ldm.common_dit
-# Copied from https://github.com/black-forest-labs/flux/blob/main/src/flux/math.py
-def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
-    assert dim % 2 == 0, "The dimension must be even."
-
-    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
-    omega = 1.0 / (theta**scale)
-
-    batch_size, seq_length = pos.shape
-    out = torch.einsum("...n,d->...nd", pos, omega)
-    cos_out = torch.cos(out)
-    sin_out = torch.sin(out)
-
-    stacked_out = torch.stack([cos_out, -sin_out, sin_out, cos_out], dim=-1)
-    out = stacked_out.view(batch_size, -1, dim // 2, 2, 2)
-    return out.float()
 
 
 # Copied from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py
@@ -84,23 +71,6 @@ class TimestepEmbed(nn.Module):
         return t_emb
 
 
-class OutEmbed(nn.Module):
-    def __init__(self, hidden_size, patch_size, out_channels, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.norm_final = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
-        self.linear = operations.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True, dtype=dtype, device=device)
-        self.adaLN_modulation = nn.Sequential(
-            nn.SiLU(),
-            operations.Linear(hidden_size, 2 * hidden_size, bias=True, dtype=dtype, device=device)
-        )
-
-    def forward(self, x, adaln_input):
-        shift, scale = self.adaLN_modulation(adaln_input).chunk(2, dim=1)
-        x = self.norm_final(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
-        x = self.linear(x)
-        return x
-
-
 def attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor):
     return optimized_attention(query.view(query.shape[0], -1, query.shape[-1] * query.shape[-2]), key.view(key.shape[0], -1, key.shape[-1] * key.shape[-2]), value.view(value.shape[0], -1, value.shape[-1] * value.shape[-2]), query.shape[2])
 
@@ -663,7 +633,7 @@ class HiDreamImageTransformer2DModel(nn.Module):
             ]
         )
 
-        self.final_layer = OutEmbed(self.inner_dim, patch_size, self.out_channels, dtype=dtype, device=device, operations=operations)
+        self.final_layer = LastLayer(self.inner_dim, patch_size, self.out_channels, dtype=dtype, device=device, operations=operations)
 
         caption_channels = [caption_channels[1], ] * (num_layers + num_single_layers) + [caption_channels[0], ]
         caption_projection = []
@@ -732,7 +702,8 @@ class HiDreamImageTransformer2DModel(nn.Module):
         control = None,
         transformer_options = {},
     ) -> torch.Tensor:
-        hidden_states = x
+        bs, c, h, w = x.shape
+        hidden_states = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size))
         timesteps = t
         pooled_embeds = y
         T5_encoder_hidden_states = context
@@ -825,4 +796,4 @@ class HiDreamImageTransformer2DModel(nn.Module):
         hidden_states = hidden_states[:, :image_tokens_seq_len, ...]
         output = self.final_layer(hidden_states, adaln_input)
         output = self.unpatchify(output, img_sizes)
-        return -output
+        return -output[:, :, :h, :w]

comfy/text_encoders/hidream.py

@@ -11,14 +11,15 @@ class HiDreamTokenizer:
     def __init__(self, embedding_directory=None, tokenizer_data={}):
         self.clip_l = sd1_clip.SDTokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)
         self.clip_g = sdxl_clip.SDXLClipGTokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)
-        self.t5xxl = sd3_clip.T5XXLTokenizer(embedding_directory=embedding_directory, min_length=128, tokenizer_data=tokenizer_data)
+        self.t5xxl = sd3_clip.T5XXLTokenizer(embedding_directory=embedding_directory, min_length=128, max_length=128, tokenizer_data=tokenizer_data)
         self.llama = hunyuan_video.LLAMA3Tokenizer(embedding_directory=embedding_directory, min_length=128, pad_token=128009, tokenizer_data=tokenizer_data)
 
     def tokenize_with_weights(self, text:str, return_word_ids=False, **kwargs):
         out = {}
         out["g"] = self.clip_g.tokenize_with_weights(text, return_word_ids)
         out["l"] = self.clip_l.tokenize_with_weights(text, return_word_ids)
-        out["t5xxl"] = self.t5xxl.tokenize_with_weights(text, return_word_ids)
+        t5xxl = self.t5xxl.tokenize_with_weights(text, return_word_ids)
+        out["t5xxl"] = [t5xxl[0]]  # Use only first 128 tokens
         out["llama"] = self.llama.tokenize_with_weights(text, return_word_ids)
         return out
 

comfy/text_encoders/sd3_clip.py

@@ -32,9 +32,9 @@ def t5_xxl_detect(state_dict, prefix=""):
     return out
 
 class T5XXLTokenizer(sd1_clip.SDTokenizer):
-    def __init__(self, embedding_directory=None, tokenizer_data={}, min_length=77):
+    def __init__(self, embedding_directory=None, tokenizer_data={}, min_length=77, max_length=99999999):
         tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "t5_tokenizer")
-        super().__init__(tokenizer_path, embedding_directory=embedding_directory, pad_with_end=False, embedding_size=4096, embedding_key='t5xxl', tokenizer_class=T5TokenizerFast, has_start_token=False, pad_to_max_length=False, max_length=99999999, min_length=min_length, tokenizer_data=tokenizer_data)
+        super().__init__(tokenizer_path, embedding_directory=embedding_directory, pad_with_end=False, embedding_size=4096, embedding_key='t5xxl', tokenizer_class=T5TokenizerFast, has_start_token=False, pad_to_max_length=False, max_length=max_length, min_length=min_length, tokenizer_data=tokenizer_data)
 
 
 class SD3Tokenizer: