wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-05-12 18:22:53 +08:00
Code Issues Actions 22 Packages Projects Releases Wiki Activity

Fix rotate node with batch_size > 1

Browse Source
This commit is contained in:
missionfloyd 2023-04-27 20:45:58 -06:00
parent 7b64c38d52
commit 03dabd2f56
1 changed files with 12 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
comfy_extras/nodes_post_processing.py
View File

@ -291,13 +291,6 @@ class Rotate:
"bicubic": Image.Resampling.BICUBIC, "bicubic": Image.Resampling.BICUBIC,
} }
tensor_image = image[0]
img = (tensor_image * 255).to(torch.uint8).numpy()
pil_image = Image.fromarray(img, mode='RGB')
expand = True if expand == "enabled" else False
fill_color = fill_color or "#000000"
def parse_palette(color_str): def parse_palette(color_str):
if re.match(r'^#[a-fA-F0-9]{6}$', color_str) or color_str.lower() in ImageColor.colormap: if re.match(r'^#[a-fA-F0-9]{6}$', color_str) or color_str.lower() in ImageColor.colormap:
return ImageColor.getrgb(color_str) return ImageColor.getrgb(color_str)
@ -308,16 +301,24 @@ class Rotate:
else: else:
raise ValueError(f"Invalid color format: {color_str}") raise ValueError(f"Invalid color format: {color_str}")
expand = True if expand == "enabled" else False
fill_color = fill_color or "#000000"
center = (width / 2, height / 2) if center_of_image == "enabled" else (center_x, center_y) center = (width / 2, height / 2) if center_of_image == "enabled" else (center_x, center_y)
translate = (translate_x, translate_y) translate = (translate_x, translate_y)
color = fill_color.replace(" ", "") color = fill_color.replace(" ", "")
color = parse_palette(color) color = parse_palette(color)
rotated_image = pil_image.rotate(angle=angle, resample=resamplers[resample], expand=expand, center=center, translate=translate, fillcolor=color)
result_width, result_height = rotated_image.size result_width, result_height = Image.new("RGB", (width, height)).rotate(angle=angle, expand=expand, center=center, translate=translate).size
result = torch.zeros(batch_size, result_height, result_width, 3) result = torch.zeros(batch_size, result_height, result_width, 3)
rotated_array = torch.tensor(np.array(rotated_image.convert("RGB"))).float() / 255
result[0] = rotated_array for b in range(batch_size):
tensor_image = image[b]
img = (tensor_image * 255).to(torch.uint8).numpy()
pil_image = Image.fromarray(img, mode='RGB')
rotated_image = pil_image.rotate(angle=angle, resample=resamplers[resample], expand=expand, center=center, translate=translate, fillcolor=color)
rotated_array = torch.tensor(np.array(rotated_image.convert("RGB"))).float() / 255
result[b] = rotated_array
return (result,) return (result,)