mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-10 06:10:50 +08:00
17 lines
754 B
Python
17 lines
754 B
Python
from typing import Final
|
|
|
|
|
|
def append_dims(x, target_dims):
|
|
"""Appends dimensions to the end of a tensor until it has target_dims dimensions."""
|
|
dims_to_append = target_dims - x.ndim
|
|
if dims_to_append < 0:
|
|
raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
|
|
expanded = x[(...,) + (None,) * dims_to_append]
|
|
# MPS will get inf values if it tries to index into the new axes, but detaching fixes this.
|
|
# https://github.com/pytorch/pytorch/issues/84364
|
|
return expanded.detach().clone() if expanded.device.type == 'mps' else expanded
|
|
|
|
|
|
class FolderOfImages:
|
|
IMG_EXTENSIONS: Final[set[str]] = frozenset({'.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp'})
|