import torch
from RealESRGAN import RealESRGAN
from PIL import Image
import os
from huggingface_hub import hf_hub_url, hf_hub_download
class Upscaler:
def __init__(self, device, scale=4, model_path=None):
"""
Initialize the Upscaler class with the specified scale factor and model
path.
:param device: The device (CUDA or CPU) where the model will be loaded.
:param scale: The scaling factor for the upscaling, either 2 or 4.
:param model_path: Path to the pre-trained model weights (optional, will
load default if None).
"""
self.device = device
self.scale = scale
self.model = None
self.model_path = model_path if model_path else
f"models/upscalers/RealESRGAN_x{self.scale}.pth"
def download_models(self):
models = {
"MODEL": ("dantea1118/juggernaut_reborn",
"juggernaut_reborn.safetensors", "models/models/Stable-diffusion"),
"UPSCALER_X2": ("ai-forever/Real-ESRGAN", "RealESRGAN_x2.pth",
"models/upscalers/"),
"UPSCALER_X4": ("ai-forever/Real-ESRGAN", "RealESRGAN_x4.pth",
"models/upscalers/"),
"NEGATIVE_1": ("philz1337x/embeddings", "verybadimagenegative_v1.3.pt",
"models/embeddings"),
"NEGATIVE_2": ("philz1337x/embeddings", "JuggernautNegative-neg.pt",
"models/embeddings"),
"LORA_1": ("philz1337x/loras", "SDXLrender_v2.0.safetensors",
"models/Lora"),
"LORA_2": ("philz1337x/loras", "more_details.safetensors",
"models/Lora"),
"CONTROLNET": ("lllyasviel/ControlNet-v1-1",
"control_v11f1e_sd15_tile.pth", "models/ControlNet"),
"VAE": ("stabilityai/sd-vae-ft-mse-original", "vae-ft-mse-840000-ema-
pruned.safetensors", "models/VAE"),
}
for model, (repo_id, filename, local_dir) in models.items():
hf_hub_download(repo_id=repo_id, filename=filename,
local_dir=local_dir)
def load_model(self):
"""
Loads the RealESRGAN model if it's not already loaded.
"""
if self.model is None:
print(f"Loading RealESRGAN model with scale {self.scale}...")
self.model = RealESRGAN(self.device, scale=self.scale)
self.model.load_weights(self.model_path, download=False)
print(f"Model with scale {self.scale} loaded successfully.")
def predict(self, img: Image) -> Image:
"""
Upscales the input image using the loaded model.
:param img: The input image to upscale.
:return: The upscaled image.
"""
self.load_model()
return self.model.predict(img)
# Example usage:
# device is determined by checking if CUDA (GPU) is available, otherwise using CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Initialize the upscaling object with the desired scale (2x or 4x)
upscaler = Upscaler(device, scale=4)
upscaler.download_models()
# Load an image to upscale
input_image = Image.open("path/to/your/image.jpg") # Replace with the path to your
image
# Perform the upscaling
upscaled_image = upscaler.predict(input_image)
# Save or display the result
upscaled_image.save("upscaled_image.jpg")
upscaled_image.show()
