71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
import torch
|
|
import numpy as np
|
|
|
|
class NearestUpscaleNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"scale_factor": ("INT", {
|
|
"default": 2,
|
|
"min": 1,
|
|
"max": 8,
|
|
"step": 1
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "upscale"
|
|
CATEGORY = "image/upscaling"
|
|
|
|
def upscale(self, image, scale_factor):
|
|
# Debug inicial do tensor de entrada
|
|
print(f"Tipo de entrada: {type(image)}")
|
|
print(f"Shape do tensor de entrada: {image.shape}")
|
|
|
|
# Converter tensor para numpy array
|
|
if isinstance(image, torch.Tensor):
|
|
image_np = image[0].cpu().numpy()
|
|
else:
|
|
image_np = np.array(image)
|
|
|
|
# Debug após conversão para numpy
|
|
print(f"Shape do numpy array: {image_np.shape}")
|
|
height, width = image_np.shape[:2]
|
|
print(f"Altura: {height}, Largura: {width}")
|
|
print(f"Fator de escala: {scale_factor}")
|
|
|
|
# Calcular novas dimensões
|
|
new_height = height * scale_factor
|
|
new_width = width * scale_factor
|
|
print(f"Nova altura calculada: {new_height}")
|
|
print(f"Nova largura calculada: {new_width}")
|
|
|
|
# Criar array de saída com as dimensões exatas
|
|
upscaled = np.zeros((new_height, new_width, 3), dtype=np.float32)
|
|
print(f"Shape do array de saída: {upscaled.shape}")
|
|
|
|
# Upscaling usando repeat nativo do numpy
|
|
# Primeiro expandimos na direção vertical
|
|
temp = np.repeat(image_np, scale_factor, axis=0)
|
|
# Depois na horizontal
|
|
upscaled = np.repeat(temp, scale_factor, axis=1)
|
|
|
|
print(f"Shape final antes do tensor: {upscaled.shape}")
|
|
|
|
# Converter para tensor mantendo as dimensões exatas
|
|
upscaled_tensor = torch.from_numpy(upscaled).unsqueeze(0)
|
|
print(f"Shape final do tensor: {upscaled_tensor.shape}")
|
|
|
|
return (upscaled_tensor,)
|
|
|
|
# Para registrar o nó
|
|
NODE_CLASS_MAPPINGS = {
|
|
"NearestUpscale": NearestUpscaleNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"NearestUpscale": "Nearest Neighbor Upscale"
|
|
} |