99from typing_extensions import override
1010
1111import folder_paths
12+ from comfy .utils import common_upscale
1213from comfy_api .latest import IO , ComfyExtension , Input
1314from comfy_api_nodes .apis .openai import (
1415 InputFileContent ,
@@ -62,7 +63,8 @@ async def validate_and_cast_response(response, timeout: int = None) -> torch.Ten
6263 timeout: Request timeout in seconds. Defaults to None (no timeout).
6364
6465 Returns:
65- A torch.Tensor representing the image (1, H, W, C).
66+ A torch.Tensor of shape (N, H, W, C) with all returned images; images whose
67+ dimensions differ from the first image's are resized to match it.
6668
6769 Raises:
6870 ValueError: If the response is not valid.
@@ -89,6 +91,14 @@ async def validate_and_cast_response(response, timeout: int = None) -> torch.Ten
8991 arr = np .asarray (pil_img ).astype (np .float32 ) / 255.0
9092 image_tensors .append (torch .from_numpy (arr ))
9193
94+ # With size="auto" the API can return images whose dimensions differ by a few pixels within a single response
95+ # resize them to the first image's dimensions so they can be stacked into one batch.
96+ ref_h , ref_w = image_tensors [0 ].shape [:2 ]
97+ for i , t in enumerate (image_tensors ):
98+ if t .shape [:2 ] != (ref_h , ref_w ):
99+ samples = t .unsqueeze (0 ).movedim (- 1 , 1 )
100+ samples = common_upscale (samples , ref_w , ref_h , "bilinear" , "center" )
101+ image_tensors [i ] = samples .movedim (1 , - 1 ).squeeze (0 )
92102 return torch .stack (image_tensors , dim = 0 )
93103
94104
0 commit comments