wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-01-31 00:30:21 +08:00
Code Issues Actions 19 Packages Projects Releases Wiki Activity

Change transpose to torch

Browse Source
This commit is contained in:
missionfloyd 2023-04-10 20:23:15 -06:00 committed by GitHub
parent 70564aebb6
commit bc54b69c59
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
comfy_extras/nodes_post_processing.py
View File

@ -236,24 +236,17 @@ class Transpose:
result = torch.permute(result, (0, 2, 1, 3)) result = torch.permute(result, (0, 2, 1, 3))
methods = { methods = {
"Flip horizontal": Image.Transpose.FLIP_LEFT_RIGHT, "Flip horizontal": (lambda x: torch.fliplr(x)),
"Flip vertical": Image.Transpose.FLIP_TOP_BOTTOM, "Flip vertical": (lambda x: torch.flipud(x)),
"Rotate 90°": Image.Transpose.ROTATE_90, "Rotate 90°": (lambda x: torch.rot90(x)),
"Rotate 180°": Image.Transpose.ROTATE_180, "Rotate 180°": (lambda x: torch.rot90(x, 2)),
"Rotate 270°": Image.Transpose.ROTATE_270, "Rotate 270°": (lambda x: torch.rot90(x, 3)),
"Transpose": Image.Transpose.TRANSPOSE, "Transpose": (lambda x: torch.transpose(x, 0, 1)),
"Transverse": Image.Transpose.TRANSVERSE, "Transverse": (lambda x: torch.rot90(torch.transpose(x, 0, 1), 2)),
} }
for b in range(batch_size): for b in range(batch_size):
tensor_image = image[b] result[b] = methods[method](image[b])
img = (tensor_image * 255).to(torch.uint8).numpy()
pil_image = Image.fromarray(img, mode='RGB')
transposed_image = pil_image.transpose(methods[method])
transposed_array = torch.tensor(np.array(transposed_image.convert("RGB"))).float() / 255
result[b] = transposed_array
return (result,) return (result,)