Run a single prediction on the model
(image, upscale, dec_w, seed, model_type, ddpm_steps, colorfix_type)
| 156 | os.makedirs('output', exist_ok=True) |
| 157 | |
| 158 | def inference(image, upscale, dec_w, seed, model_type, ddpm_steps, colorfix_type): |
| 159 | """Run a single prediction on the model""" |
| 160 | precision_scope = autocast |
| 161 | vq_model.decoder.fusion_w = dec_w |
| 162 | seed_everything(seed) |
| 163 | |
| 164 | if model_type == '512': |
| 165 | config = OmegaConf.load("./configs/stableSRNew/v2-finetune_text_T_512.yaml") |
| 166 | model = load_model_from_config(config, "./weights/stablesr_000117.ckpt") |
| 167 | min_size = 512 |
| 168 | else: |
| 169 | config = OmegaConf.load("./configs/stableSRNew/v2-finetune_text_T_768v.yaml") |
| 170 | model = load_model_from_config(config, "./weights/stablesr_768v_000139.ckpt") |
| 171 | min_size = 768 |
| 172 | |
| 173 | model = model.to(device) |
| 174 | model.configs = config |
| 175 | model.register_schedule(given_betas=None, beta_schedule="linear", timesteps=1000, |
| 176 | linear_start=0.00085, linear_end=0.0120, cosine_s=8e-3) |
| 177 | model.num_timesteps = 1000 |
| 178 | |
| 179 | sqrt_alphas_cumprod = copy.deepcopy(model.sqrt_alphas_cumprod) |
| 180 | sqrt_one_minus_alphas_cumprod = copy.deepcopy(model.sqrt_one_minus_alphas_cumprod) |
| 181 | |
| 182 | use_timesteps = set(space_timesteps(1000, [ddpm_steps])) |
| 183 | last_alpha_cumprod = 1.0 |
| 184 | new_betas = [] |
| 185 | timestep_map = [] |
| 186 | for i, alpha_cumprod in enumerate(model.alphas_cumprod): |
| 187 | if i in use_timesteps: |
| 188 | new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) |
| 189 | last_alpha_cumprod = alpha_cumprod |
| 190 | timestep_map.append(i) |
| 191 | new_betas = [beta.data.cpu().numpy() for beta in new_betas] |
| 192 | model.register_schedule(given_betas=np.array(new_betas), timesteps=len(new_betas)) |
| 193 | model.num_timesteps = 1000 |
| 194 | model.ori_timesteps = list(use_timesteps) |
| 195 | model.ori_timesteps.sort() |
| 196 | model = model.to(device) |
| 197 | |
| 198 | try: # global try |
| 199 | with torch.no_grad(): |
| 200 | with precision_scope("cuda"): |
| 201 | with model.ema_scope(): |
| 202 | init_image = load_img(image) |
| 203 | init_image = F.interpolate( |
| 204 | init_image, |
| 205 | size=(int(init_image.size(-2)*upscale), |
| 206 | int(init_image.size(-1)*upscale)), |
| 207 | mode='bicubic', |
| 208 | ) |
| 209 | |
| 210 | if init_image.size(-1) < min_size or init_image.size(-2) < min_size: |
| 211 | ori_size = init_image.size() |
| 212 | rescale = min_size * 1.0 / min(init_image.size(-2), init_image.size(-1)) |
| 213 | new_h = max(int(ori_size[-2]*rescale), min_size) |
| 214 | new_w = max(int(ori_size[-1]*rescale), min_size) |
| 215 | init_template = F.interpolate( |
nothing calls this directly
no test coverage detected