37 lines
1.1 KiB
Python
37 lines
1.1 KiB
Python
import numpy as np
|
|
from PIL import Image
|
|
import torch
|
|
import random
|
|
import gc
|
|
from comfy.model_management import unload_all_models, soft_empty_cache
|
|
|
|
def tensor2pil(image):
|
|
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
|
|
|
def pil2tensor(image):
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
def get_image_size(image_path):
|
|
# 打开图像文件
|
|
with Image.open(image_path) as img:
|
|
# 获取图像的宽度和高度
|
|
width, height = img.size
|
|
return width, height
|
|
|
|
def clear_memory():
|
|
gc.collect()
|
|
unload_all_models()
|
|
soft_empty_cache(True)
|
|
|
|
def set_seed(seed):
|
|
"""Set random seeds for reproducibility"""
|
|
if seed > 0:
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
torch.backends.cudnn.deterministic = True
|
|
torch.backends.cudnn.benchmark = False
|
|
return seed |