File size: 2,517 Bytes
6a6edcb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import os
import yaml
import torch
import torchvision
from tqdm import tqdm
from inference.utils import *
from train import ControlNetCore, WurstCoreB
import warnings
warnings.filterwarnings("ignore")
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
class Upscale_CaseCade:
def __init__(self) -> None:
self.config_file = './configs/inference/controlnet_c_3b_sr.yaml'
# SETUP STAGE C
with open(self.config_file, "r", encoding="utf-8") as file:
loaded_config = yaml.safe_load(file)
self.core = ControlNetCore(config_dict=loaded_config, device=device, training=False)
# SETUP STAGE B
self.config_file_b = './configs/inference/stage_b_3b.yaml'
with open(self.config_file_b, "r", encoding="utf-8") as file:
self.config_file_b = yaml.safe_load(file)
self.core_b = WurstCoreB(config_dict=self.config_file_b, device=device, training=False)
self.extras = self.core.setup_extras_pre()
self.models = self.core.setup_models(self.extras)
self.models.generator.eval().requires_grad_(False)
print("CONTROLNET READY")
self.extras_b = self.core_b.setup_extras_pre()
self.models_b = self.core_b.setup_models(self.extras_b, skip_clip=True)
self.models_b = WurstCoreB.Models(
**{**self.models_b.to_dict(), 'tokenizer': self.models.tokenizer, 'text_model': self.models.text_model}
)
self.models_b.generator.eval().requires_grad_(False)
print("STAGE B READY")
def upscale_image(self,image_pil,scale_fator):
batch_size = 1
cnet_override = None
images = resize_image(image_pil).unsqueeze(0).expand(batch_size, -1, -1, -1)
batch = {'images': images}
with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16):
effnet_latents = self.core.encode_latents(batch, self.models, self.extras)
effnet_latents_up = torch.nn.functional.interpolate(effnet_latents, scale_factor=scale_fator, mode="nearest")
cnet = self.models.controlnet(effnet_latents_up)
cnet_uncond = cnet
cnet_input = torch.nn.functional.interpolate(images, scale_factor=scale_fator, mode="nearest")
# cnet, cnet_input = core.get_cnet(batch, models, extras)
# cnet_uncond = cnet
og=show_images(batch['images'],return_images=True)
upsclae=show_images(cnet_input,return_images=True)
return og,upsclae
|