mirror of
https://github.com/vladmandic/sdnext.git
synced 2026-01-27 15:02:48 +03:00
57 lines
1.5 KiB
Python
57 lines
1.5 KiB
Python
import os
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from torchvision.transforms import ToPILImage
|
|
from .src.core.generation import generation_loop
|
|
from .src.core.model_manager import configure_runner
|
|
|
|
|
|
|
|
device = 'cuda'
|
|
dtype = torch.bfloat16
|
|
model_dir = 'seedvr2_models'
|
|
model = 'seedvr2_ema_3b_fp16.safetensors'
|
|
resolution = 1024
|
|
seed = 100
|
|
cfg = 1.0
|
|
input_image = '/home/vlado/generative/Samples/cutie-512.png'
|
|
|
|
to_pil = ToPILImage()
|
|
runner = None
|
|
loaded_model = None
|
|
|
|
|
|
def upscale_image(model_name:str, image_path:str):
|
|
global runner, loaded_model
|
|
if (runner is None) or (loaded_model != model_name):
|
|
runner = configure_runner(model_name, model_dir, device=device, dtype=dtype)
|
|
loaded_model = model_name
|
|
|
|
image = Image.open(image_path).convert("RGB")
|
|
image_tensor = np.array(image)
|
|
image_tensor = torch.from_numpy(image_tensor).to(device=device, dtype=dtype).unsqueeze(0) / 255.0
|
|
|
|
result_tensor = generation_loop(
|
|
runner=runner,
|
|
images=image_tensor,
|
|
cfg_scale=cfg,
|
|
seed=seed,
|
|
res_w=resolution,
|
|
batch_size=1,
|
|
temporal_overlap=0,
|
|
device=device,
|
|
)
|
|
image = to_pil(result_tensor.squeeze().permute((2, 0, 1)))
|
|
|
|
output_path = os.path.join('/tmp', os.path.basename(image_path))
|
|
|
|
image.save(output_path, quality=95)
|
|
return image
|
|
|
|
|
|
if __name__ == "__main__":
|
|
output_image = upscale_image(model, input_image)
|
|
print('input:', input_image)
|
|
print('output:', output_image)
|