mirror of
https://github.com/Nighthawk42/manga-colorizer.git
synced 2026-08-30 07:22:28 +00:00
20 lines
518 B
Python
20 lines
518 B
Python
import torch
|
|
|
|
from denoising.denoiser import FFDNetDenoiser
|
|
|
|
|
|
class MangaDenoiser:
|
|
def __init__(self, config):
|
|
if config.device == 'cuda' and not torch.cuda.is_available():
|
|
print("[-] CUDA not available, using CPU.")
|
|
self.device = 'cpu'
|
|
else:
|
|
self.device = config.device
|
|
|
|
self.model = FFDNetDenoiser(self.device)
|
|
|
|
|
|
def denoise(self, image, sigma=25):
|
|
with torch.no_grad():
|
|
return self.model.get_denoised_image(image, sigma=sigma)
|