wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-01-11 06:40:48 +08:00
Code Issues Actions 12 Packages Projects Releases Wiki Activity

Fix tensor types

Browse Source
This commit is contained in:
doctorpangloss 2024-09-06 11:04:32 -07:00
parent a4fb34a0b8
commit e8eab4dbc6
1 changed files with 11 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
comfy/component_model/tensor_types.py
View File

@ -1,8 +1,19 @@
from typing import TypedDict
from jaxtyping import Float
from torch import Tensor
from typing_extensions import NotRequired
ImageBatch = Float[Tensor, "batch height width channels"]
LatentBatch = Float[Tensor, "batch channels height width"]
SD1LatentBatch = Float[Tensor, "batch 8 height width"]
SD3LatentBatch = Float[Tensor, "batch 16 height width"]
MaskBatch = Float[Tensor, "batch height width"]
RGBImageBatch = Float[Tensor, "batch height width 3"]
RGBAImageBatch = Float[Tensor, "batch height width 4"]
RGBImage = Float[Tensor, "height width 3"]
class Latent(TypedDict):
samples: LatentBatch
noise_mask: NotRequired[LatentBatch]