39 lines
1.0 KiB
Python
39 lines
1.0 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Sequence
|
|
|
|
import numpy as np
|
|
import torch
|
|
from comfy import model_management
|
|
from PIL import Image as PILImage
|
|
|
|
|
|
def pilimgbatch_to_torch(
|
|
imgbatch: Sequence[PILImage, ...] | torch.Tensor,
|
|
) -> torch.Tensor:
|
|
if isinstance(imgbatch, torch.Tensor):
|
|
print("pibtt: skip", imgbatch.shape, imgbatch.min(), imgbatch.max())
|
|
return imgbatch
|
|
npi = np.stack(
|
|
tuple(np.array(i).astype(np.float32) / 255.0 for i in imgbatch),
|
|
axis=0,
|
|
)
|
|
return torch.from_numpy(npi)
|
|
return torch.from_numpy(npi.transpose(0, 3, 1, 2))
|
|
|
|
|
|
def torch_to_pilimgbatch(t: torch.Tensor) -> tuple[PILImage, ...]:
|
|
return tuple(
|
|
PILImage.fromarray(
|
|
np.clip((255.0 * i).cpu().numpy(), 0, 255).astype(np.uint8),
|
|
)
|
|
for i in t
|
|
)
|
|
|
|
|
|
def ensure_model(model):
|
|
mp = model.inner_model.model_patcher
|
|
if model_management.LoadedModel(mp) in model_management.current_loaded_models:
|
|
return
|
|
model_management.load_models_gpu((mp,))
|