diff --git a/comfy/ldm/krea2/model.py b/comfy/ldm/krea2/model.py index ecb16254fae..8001812d77d 100644 --- a/comfy/ldm/krea2/model.py +++ b/comfy/ldm/krea2/model.py @@ -15,6 +15,7 @@ import comfy.model_management import comfy.patcher_extension import comfy.ldm.common_dit +import comfy.utils from comfy.ldm.flux.layers import EmbedND, timestep_embedding from comfy.ldm.flux.math import apply_rope from comfy.ldm.modules.attention import optimized_attention_masked @@ -73,11 +74,20 @@ def __init__(self, dim: int, heads: int, kvheads: Optional[int] = None, bias: bo self.wo = operations.Linear(dim, dim, bias=bias, device=device, dtype=dtype) def forward(self, x, freqs=None, mask=None, transformer_options={}): + transformer_patches = transformer_options.get("patches", {}) + extra_options = transformer_options.copy() q, k, v, gate = self.wq(x), self.wk(x), self.wv(x), self.gate(x) q = rearrange(q, "B L (H D) -> B H L D", H=self.heads) k = rearrange(k, "B L (H D) -> B H L D", H=self.kvheads) v = rearrange(v, "B L (H D) -> B H L D", H=self.kvheads) q, k = self.qknorm(q, k) + + if "block_index" in transformer_options and "attn1_patch" in transformer_patches: + for p in transformer_patches["attn1_patch"]: + out = p(q, k, v, pe=freqs, attn_mask=mask, extra_options=extra_options) + q, k, v = out.get("q", q), out.get("k", k), out.get("v", v) + freqs, mask = out.get("pe", freqs), out.get("attn_mask", mask) + if freqs is not None: q, k = apply_rope(q, k, freqs) if self.kvheads != self.heads: @@ -86,6 +96,11 @@ def forward(self, x, freqs=None, mask=None, transformer_options={}): v = v.repeat_interleave(rep, dim=1) out = optimized_attention_masked(q, k, v, self.heads, mask=mask, skip_reshape=True, transformer_options=transformer_options) + + if "block_index" in transformer_options and "attn1_output_patch" in transformer_patches: + for p in transformer_patches["attn1_output_patch"]: + out = p(out, extra_options) + return self.wo(out * F.sigmoid(gate)) @@ -158,8 +173,44 @@ def __init__(self, features, heads, multiplier, bias=False, kvheads=None, device self.attn = Attention(features, heads, kvheads=kvheads, bias=bias, device=device, dtype=dtype, operations=operations) self.mlp = SwiGLU(features, multiplier, bias, device=device, dtype=dtype, operations=operations) - def forward(self, x, vec, freqs, mask=None, transformer_options={}): + def forward(self, x, vec, freqs, mask=None, timestep_zero_index=None, transformer_options={}): prescale, preshift, pregate, postscale, postshift, postgate = self.mod(vec) + if timestep_zero_index is not None: + bs = x.shape[0] + ref_prescale = prescale[bs:] + ref_preshift = preshift[bs:] + ref_pregate = pregate[bs:] + ref_postscale = postscale[bs:] + ref_postshift = postshift[bs:] + ref_postgate = postgate[bs:] + prescale = prescale[:bs] + preshift = preshift[:bs] + pregate = pregate[:bs] + postscale = postscale[:bs] + postshift = postshift[:bs] + postgate = postgate[:bs] + + pre = self.prenorm(x) + pre[:, :timestep_zero_index].mul_(1 + prescale).add_(preshift) + pre[:, timestep_zero_index:].mul_(1 + ref_prescale).add_(ref_preshift) + attn = self.attn(pre, freqs, mask, transformer_options=transformer_options) + del pre + attn[:, :timestep_zero_index].mul_(pregate) + attn[:, timestep_zero_index:].mul_(ref_pregate) + x = x + attn + del attn + + post = self.postnorm(x) + post[:, :timestep_zero_index].mul_(1 + postscale).add_(postshift) + post[:, timestep_zero_index:].mul_(1 + ref_postscale).add_(ref_postshift) + mlp = self.mlp(post) + del post + mlp[:, :timestep_zero_index].mul_(postgate) + mlp[:, timestep_zero_index:].mul_(ref_postgate) + x = x + mlp + del mlp + return x + x = x + pregate * self.attn((1 + prescale) * self.prenorm(x) + preshift, freqs, mask, transformer_options=transformer_options) x = x + postgate * self.mlp((1 + postscale) * self.postnorm(x) + postshift) return x @@ -181,7 +232,7 @@ def forward(self, x, tvec): class SingleStreamDiT(nn.Module): def __init__(self, features=6144, tdim=256, txtdim=2560, heads=48, kvheads=12, multiplier=4, layers=28, patch=2, channels=16, bias=False, theta=1e3, txtlayers=12, - txtheads=20, txtkvheads=20, image_model=None, + txtheads=20, txtkvheads=20, default_ref_method=None, image_model=None, device=None, dtype=None, operations=None, **kwargs): super().__init__() self.dtype = dtype @@ -191,6 +242,7 @@ def __init__(self, features=6144, tdim=256, txtdim=2560, heads=48, kvheads=12, m self.heads = heads self.txtdim = txtdim self.txtlayers = txtlayers + self.default_ref_method = default_ref_method headdim = features // heads axes = [headdim - 12 * (headdim // 16), 6 * (headdim // 16), 6 * (headdim // 16)] @@ -221,61 +273,110 @@ def __init__(self, features=6144, tdim=256, txtdim=2560, heads=48, kvheads=12, m operations.Linear(features, features * 6, device=device, dtype=dtype), ) - def forward(self, x, timesteps, context, attention_mask=None, transformer_options={}, **kwargs): + def forward(self, x, timesteps, context, attention_mask=None, ref_latents=None, transformer_options={}, **kwargs): return comfy.patcher_extension.WrapperExecutor.new_class_executor( self._forward, self, comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options), - ).execute(x, timesteps, context, attention_mask, transformer_options, **kwargs) + ).execute(x, timesteps, context, attention_mask, ref_latents, transformer_options, **kwargs) + + def process_img(self, x, index=0): + patch = self.patch + x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch)) + h, w = x.shape[-2] // patch, x.shape[-1] // patch + img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch) - def _forward(self, x, timesteps, context, attention_mask=None, transformer_options={}, **kwargs): + img_ids = torch.zeros(h, w, 3, device=x.device, dtype=torch.float32) + img_ids[..., 0] = index + img_ids[..., 1] = torch.arange(h, device=x.device, dtype=torch.float32)[:, None] + img_ids[..., 2] = torch.arange(w, device=x.device, dtype=torch.float32)[None, :] + return img, img_ids.reshape(1, h * w, 3).repeat(x.shape[0], 1, 1), h, w + + def _forward(self, x, timesteps, context, attention_mask=None, ref_latents=None, transformer_options={}, **kwargs): + transformer_options = transformer_options.copy() temporal = x.ndim == 5 if temporal: b5, c5, t5, h5, w5 = x.shape x = x.reshape(b5 * t5, c5, h5, w5) - bs, c, H_orig, W_orig = x.shape + bs, _, h_orig, w_orig = x.shape patch = self.patch - # Pad the latent up to a multiple of patch (as Flux/Lumina/QwenImage do); crop back at the end. - x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch)) - H, W = x.shape[-2], x.shape[-1] - h_, w_ = H // patch, W // patch # context arrives as (B, seq, txtlayers*txtdim); reshape to (B, txtlayers, seq, txtdim). context = self._unpack_context(context) - img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch) + img, imgpos, h_, w_ = self.process_img(x) + img_tokens = img.shape[1] + timestep_zero_index = None + ref_method = kwargs.get("ref_latents_method", self.default_ref_method) + if ref_method is not None and ref_latents is not None and len(ref_latents) > 0: + ref_tokens = [] + ref_pos = [] + ref_num_tokens = [] + for index, ref in enumerate(ref_latents, 1): + if ref.ndim == 5: + rb, rc, rt, rh5, rw5 = ref.shape + ref = ref.reshape(rb * rt, rc, rh5, rw5) + ref = comfy.utils.repeat_to_batch_size(ref, bs) + kontext, kontext_ids, _, _ = self.process_img(ref, index=index) + ref_tokens.append(kontext) + ref_pos.append(kontext_ids) + ref_num_tokens.append(kontext.shape[1]) + img = torch.cat([img] + ref_tokens, dim=1) + imgpos = torch.cat([imgpos] + ref_pos, dim=1) + del ref_tokens, ref_pos + if ref_method == "index_timestep_zero": + timestep_zero_index = img_tokens + transformer_options["reference_image_num_tokens"] = ref_num_tokens + img = self.first(img) t = self.tmlp(timestep_embedding(timesteps, self.tdim).unsqueeze(1).to(img.dtype)) tvec = self.tproj(t) + if timestep_zero_index is not None: + t0 = self.tmlp(timestep_embedding(torch.zeros_like(timesteps), self.tdim).unsqueeze(1).to(img.dtype)) + tvec = torch.cat((tvec, self.tproj(t0)), dim=0) context = self.txtfusion(context, mask=None, transformer_options=transformer_options) context = self.txtmlp(context) - txtlen, imglen = context.shape[1], img.shape[1] + txtlen = context.shape[1] + device = context.device + txtpos = torch.zeros(bs, txtlen, 3, device=device, dtype=torch.float32) + + patches = transformer_options.get("patches", {}) + if "post_input" in patches: + for p in patches["post_input"]: + out = p({"img": img, "txt": context, "img_ids": imgpos, "txt_ids": txtpos, "transformer_options": transformer_options}) + img, context = out["img"], out["txt"] + imgpos, txtpos = out["img_ids"], out["txt_ids"] + combined = torch.cat((context, img), dim=1) + del context, img + if timestep_zero_index is not None: + timestep_zero_index += txtlen # Position ids: text at 0, image at (0, h_idx, w_idx). - device = combined.device - txtpos = torch.zeros(bs, txtlen, 3, device=device, dtype=torch.float32) - imgids = torch.zeros(h_, w_, 3, device=device, dtype=torch.float32) - imgids[..., 1] = torch.arange(h_, device=device, dtype=torch.float32)[:, None] - imgids[..., 2] = torch.arange(w_, device=device, dtype=torch.float32)[None, :] - imgpos = imgids.reshape(1, h_ * w_, 3).repeat(bs, 1, 1) pos = torch.cat((txtpos, imgpos), dim=1) + del txtpos, imgpos freqs = self.pe_embedder(pos) + del pos - for block in self.blocks: - combined = block(combined, tvec, freqs, None, transformer_options=transformer_options) + transformer_options["total_blocks"] = len(self.blocks) + transformer_options["block_type"] = "single" + transformer_options["img_slice"] = [txtlen, combined.shape[1]] + for i, block in enumerate(self.blocks): + transformer_options["block_index"] = i + combined = block(combined, tvec, freqs, None, timestep_zero_index=timestep_zero_index, transformer_options=transformer_options) final = self.last(combined, t) - out = final[:, txtlen:txtlen + imglen, :] + del combined + out = final[:, txtlen:txtlen + img_tokens, :] out = rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_, w=w_, ph=patch, pw=patch, c=self.channels) - out = out[:, :, :H_orig, :W_orig] # crop padding back off + out = out[:, :, :h_orig, :w_orig] # crop padding back off if temporal: - out = out.reshape(b5, t5, self.channels, H_orig, W_orig).movedim(1, 2) + out = out.reshape(b5, t5, self.channels, h_orig, w_orig).movedim(1, 2) return out def _unpack_context(self, context): diff --git a/comfy/model_base.py b/comfy/model_base.py index 98f5ba48b59..0f705316cd1 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -2227,10 +2227,7 @@ def extra_conds(self, **kwargs): 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) + out['ref_latents'] = comfy.conds.CONDList([self.process_latent_in(lat) for lat in ref_latents]) return out def extra_conds_shapes(self, **kwargs): @@ -2317,12 +2314,30 @@ def extra_conds(self, **kwargs): class Krea2(BaseModel): def __init__(self, model_config, model_type=ModelType.FLUX, device=None): super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.krea2.model.SingleStreamDiT) + self.memory_usage_factor_conds = ("ref_latents",) def extra_conds(self, **kwargs): out = super().extra_conds(**kwargs) 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, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16]) return out class HunyuanImage21(BaseModel): diff --git a/comfy_api_nodes/apis/gemini.py b/comfy_api_nodes/apis/gemini.py index 7b25432708f..ce89928cb9b 100644 --- a/comfy_api_nodes/apis/gemini.py +++ b/comfy_api_nodes/apis/gemini.py @@ -1,6 +1,6 @@ from datetime import date from enum import Enum -from typing import Any +from typing import Any, Literal from pydantic import BaseModel, Field @@ -242,3 +242,60 @@ class GeminiGenerateContentResponse(BaseModel): promptFeedback: GeminiPromptFeedback | None = Field(None) usageMetadata: GeminiUsageMetadata | None = Field(None) modelVersion: str | None = Field(None) + + +class GeminiInteractionTextPart(BaseModel): + type: Literal["text"] = "text" + text: str = Field(...) + + +class GeminiInteractionMediaPart(BaseModel): + type: str = Field(..., description="One of: image, video, audio, document.") + data: str | None = Field(None, description="Base64-encoded media bytes.") + uri: str | None = Field(None, description="URI of the media, as an alternative to inline data.") + mime_type: str | None = Field(None) + + +class GeminiInteractionGenerationConfig(BaseModel): + temperature: float | None = Field(None, ge=0.0, le=2.0) + top_p: float | None = Field(None, ge=0.0, le=1.0) + + +class GeminiInteractionRequest(BaseModel): + model: str = Field(...) + input: list[GeminiInteractionTextPart | GeminiInteractionMediaPart] = Field(...) + generation_config: GeminiInteractionGenerationConfig | None = Field(None) + + +class GeminiInteractionModalityTokens(BaseModel): + modality: str | None = Field(None, description="One of: text, image, audio, video, document.") + tokens: int | None = Field(None) + + +class GeminiInteractionUsage(BaseModel): + input_tokens_by_modality: list[GeminiInteractionModalityTokens] | None = Field(None) + output_tokens_by_modality: list[GeminiInteractionModalityTokens] | None = Field(None) + total_thought_tokens: int | None = Field(None) + + +class GeminiInteractionContent(BaseModel): + type: str | None = Field(None) + text: str | None = Field(None) + data: str | None = Field(None) + uri: str | None = Field(None) + mime_type: str | None = Field(None) + + +class GeminiInteractionStep(BaseModel): + type: str | None = Field(None) + content: list[GeminiInteractionContent] | None = Field(None) + + +class GeminiInteraction(BaseModel): + id: str | None = Field(None) + status: str | None = Field( + None, + description="One of: in_progress, requires_action, completed, failed, cancelled, incomplete.", + ) + steps: list[GeminiInteractionStep] | None = Field(None) + usage: GeminiInteractionUsage | None = Field(None) diff --git a/comfy_api_nodes/nodes_gemini.py b/comfy_api_nodes/nodes_gemini.py index 283e2233df8..8998b5943e9 100644 --- a/comfy_api_nodes/nodes_gemini.py +++ b/comfy_api_nodes/nodes_gemini.py @@ -24,6 +24,11 @@ GeminiImageGenerateContentRequest, GeminiImageGenerationConfig, GeminiInlineData, + GeminiInteraction, + GeminiInteractionGenerationConfig, + GeminiInteractionMediaPart, + GeminiInteractionRequest, + GeminiInteractionTextPart, GeminiMimeType, GeminiPart, GeminiRole, @@ -51,6 +56,7 @@ ) GEMINI_BASE_ENDPOINT = "/proxy/vertexai/gemini" +GEMINI_INTERACTIONS_ENDPOINT = "/proxy/gemini-interactions" GEMINI_MAX_INPUT_FILE_SIZE = 20 * 1024 * 1024 # 20 MB GEMINI_URL_INPUT_BUDGET = 10 GEMINI_MAX_INLINE_BYTES = 18 * 1024 * 1024 @@ -231,29 +237,10 @@ async def get_image_from_response(response: GeminiGenerateContentResponse, thoug return torch.cat(image_tensors, dim=0) -async def get_video_from_response( - response: GeminiGenerateContentResponse, cls: type[IO.ComfyNode] | None = None -) -> InputImpl.VideoFromFile: - parts = get_parts_by_type(response, "video/*") - for part in parts: - if part.inlineData and part.inlineData.data: - return InputImpl.VideoFromFile(BytesIO(base64.b64decode(part.inlineData.data))) - if part.fileData and part.fileData.fileUri: - return await download_url_to_video_output(part.fileData.fileUri, cls=cls) - model_message = get_text_from_response(response).strip() - if model_message: - raise ValueError(f"Gemini did not generate a video. Model response: {model_message}") - raise ValueError( - "Gemini did not generate a video. Try rephrasing your prompt, " - "shortening the requested duration, or reducing the number of input images/videos." - ) - - def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | None: if not response.modelVersion: return None # Define prices (Cost per 1,000,000 tokens), see https://cloud.google.com/vertex-ai/generative-ai/pricing - output_video_tokens_price = 0.0 if response.modelVersion == "gemini-2.5-pro": input_tokens_price = 1.25 output_text_tokens_price = 10.0 @@ -290,11 +277,6 @@ def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | N input_tokens_price = 0.25 output_text_tokens_price = 1.50 output_image_tokens_price = 30.0 - elif response.modelVersion == "gemini-omni-flash-preview": - input_tokens_price = 2.145 - output_text_tokens_price = 12.87 - output_image_tokens_price = 0.0 - output_video_tokens_price = 25.025 else: return None final_price = response.usageMetadata.promptTokenCount * input_tokens_price @@ -302,8 +284,6 @@ def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | N for i in response.usageMetadata.candidatesTokensDetails: if i.modality == Modality.IMAGE: final_price += output_image_tokens_price * i.tokenCount # for Nano Banana models - elif i.modality == Modality.VIDEO: - final_price += output_video_tokens_price * i.tokenCount # for Omni Flash else: final_price += output_text_tokens_price * i.tokenCount if response.usageMetadata.thoughtsTokenCount: @@ -311,6 +291,58 @@ def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | N return final_price / 1_000_000.0 +def get_text_from_interaction(interaction: GeminiInteraction) -> str: + """Extract and concatenate all model output text from an Interactions API response.""" + texts = [] + for step in interaction.steps or []: + if step.type != "model_output": + continue + for content in step.content or []: + if content.type == "text" and content.text: + texts.append(content.text) + return "\n".join(texts) + + +async def get_video_from_interaction( + interaction: GeminiInteraction, cls: type[IO.ComfyNode] | None = None +) -> InputImpl.VideoFromFile: + for step in interaction.steps or []: + if step.type != "model_output": + continue + for content in step.content or []: + if content.type != "video": + continue + if content.data: + return InputImpl.VideoFromFile(BytesIO(base64.b64decode(content.data))) + if content.uri: + return await download_url_to_video_output(content.uri, cls=cls) + model_message = get_text_from_interaction(interaction).strip() + if model_message: + raise ValueError(f"Gemini did not generate a video. Model response: {model_message}") + raise ValueError( + "Gemini did not generate a video. Try rephrasing your prompt, " + "shortening the requested duration, or reducing the number of input images/videos." + ) + + +def calculate_interaction_tokens_price(interaction: GeminiInteraction) -> float | None: + if interaction.usage is None: + return None + input_tokens_price = 1.5 + output_tokens_prices = {"text": 9.0, "video": 17.5} + thoughts_tokens_price = 9.0 + final_price = 0.0 + for i in interaction.usage.input_tokens_by_modality or []: + if i.tokens: + final_price += input_tokens_price * i.tokens + for i in interaction.usage.output_tokens_by_modality or []: + if i.tokens and i.modality in output_tokens_prices: + final_price += output_tokens_prices[i.modality] * i.tokens + if interaction.usage.total_thought_tokens: + final_price += thoughts_tokens_price * interaction.usage.total_thought_tokens + return final_price / 1_000_000.0 + + def create_video_parts(video_input: Input.Video) -> list[GeminiPart]: """Convert a single video input to Gemini API compatible parts (inline MP4/H.264).""" base_64_string = video_to_base64_string( @@ -445,6 +477,15 @@ async def build_gemini_media_parts( return parts +def to_interaction_media_part(part: GeminiPart) -> GeminiInteractionMediaPart: + """Convert a fileData/inlineData GeminiPart into an Interactions API media part.""" + if part.fileData: + mime = part.fileData.mimeType.value + return GeminiInteractionMediaPart(type=mime.split("/")[0], uri=part.fileData.fileUri, mime_type=mime) + mime = part.inlineData.mimeType.value + return GeminiInteractionMediaPart(type=mime.split("/")[0], data=part.inlineData.data, mime_type=mime) + + class GeminiNode(IO.ComfyNode): """ Node to generate text responses from a Gemini model. @@ -1676,7 +1717,7 @@ def define_schema(cls): ], is_api_node=True, price_badge=IO.PriceBadge( - expr='{"type":"usd","usd":0.146,"format":{"suffix":"/second","approximate":true}}' + expr='{"type":"usd","usd":0.101,"format":{"suffix":"/second","approximate":true}}' ), ) @@ -1695,27 +1736,34 @@ async def execute(cls, model: dict, seed: int) -> IO.NodeOutput: for video in videos: validate_video_duration(video, max_duration=10) - parts: list[GeminiPart] = [] + parts: list[GeminiInteractionTextPart | GeminiInteractionMediaPart] = [] if images or videos: - parts.extend(await build_gemini_media_parts(cls, images, [], videos)) - parts.append(GeminiPart(text=prompt)) - response = await sync_op( + media_parts = await build_gemini_media_parts(cls, images, [], videos) + parts.extend(to_interaction_media_part(p) for p in media_parts) + parts.append(GeminiInteractionTextPart(text=prompt)) + interaction = await sync_op( cls, - ApiEndpoint(path=f"{GEMINI_BASE_ENDPOINT}/{model_id}", method="POST"), - data=GeminiGenerateContentRequest( - contents=[GeminiContent(role=GeminiRole.user, parts=parts)], - generationConfig=GeminiGenerationConfig( - responseModalities=["TEXT", "VIDEO"], + ApiEndpoint(path=GEMINI_INTERACTIONS_ENDPOINT, method="POST"), + data=GeminiInteractionRequest( + model=model_id, + input=parts, + generation_config=GeminiInteractionGenerationConfig( temperature=model.get("temperature", 1.0), - topP=model.get("top_p", 0.95), + top_p=model.get("top_p", 0.95), ), ), - response_model=GeminiGenerateContentResponse, - price_extractor=calculate_tokens_price, + response_model=GeminiInteraction, + price_extractor=calculate_interaction_tokens_price, ) + if interaction.status != "completed": + model_message = get_text_from_interaction(interaction).strip() + raise ValueError( + f"Gemini interaction did not complete (status: {interaction.status})." + + (f" Model response: {model_message}" if model_message else "") + ) return IO.NodeOutput( - await get_video_from_response(response, cls=cls), - get_text_from_response(response), + await get_video_from_interaction(interaction, cls=cls), + get_text_from_interaction(interaction), )