Run a single prediction on the model
(
self,
input_image: Path = Input(description="Input image"),
ddpm_steps: int = Input(
description="Number of DDPM steps for sampling", default=200
),
fidelity_weight: float = Input(
description="Balance the quality (lower number) and fidelity (higher number)",
default=0.5,
),
upscale: float = Input(
description="The upscale for super-resolution, 4x SR by default",
default=4.0,
),
tile_overlap: int = Input(
description="The overlap between tiles, betwwen 0 to 64",
ge=0,
le=64,
default=32,
),
colorfix_type: str = Input(
choices=["adain", "wavelet", "none"], default="adain"
),
seed: int = Input(
description="Random seed. Leave blank to randomize the seed", default=None
),
)
| 41 | self.vq_model = self.vq_model.to(device) |
| 42 | |
| 43 | def predict( |
| 44 | self, |
| 45 | input_image: Path = Input(description="Input image"), |
| 46 | ddpm_steps: int = Input( |
| 47 | description="Number of DDPM steps for sampling", default=200 |
| 48 | ), |
| 49 | fidelity_weight: float = Input( |
| 50 | description="Balance the quality (lower number) and fidelity (higher number)", |
| 51 | default=0.5, |
| 52 | ), |
| 53 | upscale: float = Input( |
| 54 | description="The upscale for super-resolution, 4x SR by default", |
| 55 | default=4.0, |
| 56 | ), |
| 57 | tile_overlap: int = Input( |
| 58 | description="The overlap between tiles, betwwen 0 to 64", |
| 59 | ge=0, |
| 60 | le=64, |
| 61 | default=32, |
| 62 | ), |
| 63 | colorfix_type: str = Input( |
| 64 | choices=["adain", "wavelet", "none"], default="adain" |
| 65 | ), |
| 66 | seed: int = Input( |
| 67 | description="Random seed. Leave blank to randomize the seed", default=None |
| 68 | ), |
| 69 | ) -> Path: |
| 70 | """Run a single prediction on the model""" |
| 71 | if seed is None: |
| 72 | seed = int.from_bytes(os.urandom(2), "big") |
| 73 | print(f"Using seed: {seed}") |
| 74 | |
| 75 | self.vq_model.decoder.fusion_w = fidelity_weight |
| 76 | |
| 77 | seed_everything(seed) |
| 78 | |
| 79 | n_samples = 1 |
| 80 | device = torch.device("cuda") |
| 81 | |
| 82 | cur_image = load_img(str(input_image)).to(device) |
| 83 | cur_image = F.interpolate( |
| 84 | cur_image, |
| 85 | size=(int(cur_image.size(-2) * upscale), int(cur_image.size(-1) * upscale)), |
| 86 | mode="bicubic", |
| 87 | ) |
| 88 | |
| 89 | self.model.register_schedule( |
| 90 | given_betas=None, |
| 91 | beta_schedule="linear", |
| 92 | timesteps=1000, |
| 93 | linear_start=0.00085, |
| 94 | linear_end=0.0120, |
| 95 | cosine_s=8e-3, |
| 96 | ) |
| 97 | self.model.num_timesteps = 1000 |
| 98 | |
| 99 | sqrt_alphas_cumprod = copy.deepcopy(self.model.sqrt_alphas_cumprod) |
| 100 | sqrt_one_minus_alphas_cumprod = copy.deepcopy( |
nothing calls this directly
no test coverage detected