From 4a9010165dd85f30b59209f88033fcd3bc9205b7 Mon Sep 17 00:00:00 2001 From: edwin-yan Date: Tue, 17 Oct 2023 23:33:33 -0500 Subject: [PATCH] Added MPS for Mac Users --- api.py | 47 +++++++++++++++++++++++++++++++---------------- 1 file changed, 31 insertions(+), 16 deletions(-) diff --git a/api.py b/api.py index 1b7cf56..22c0e33 100644 --- a/api.py +++ b/api.py @@ -3,10 +3,25 @@ from helper import * from model.generator import SkipEncoderDecoder, input_noise -def remove_watermark(image_path, mask_path, max_dim, reg_noise, input_depth, lr, show_step, training_steps, tqdm_length = 100): - DTYPE = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor - if not torch.cuda.is_available(): - print('\nSetting device to "cpu", since torch is not built with "cuda" support...') + +def remove_watermark(image_path, mask_path, max_dim, reg_noise, input_depth, lr, show_step, training_steps, tqdm_length=100): + DTYPE = torch.FloatTensor + has_set_device = False + if torch.cuda.is_available(): + device = 'cuda' + has_set_device = True + print("Setting Device to CUDA...") + try: + if torch.backends.mps.is_available(): + device = 'mps' + has_set_device = True + print("Setting Device to MPS...") + except Exception as e: + print(f"Your version of pytorch might be too old, which does not support MPS. Error: \n{e}") + pass + if not has_set_device: + device = 'cpu' + print('\nSetting device to "cpu", since torch is not built with "cuda" or "mps" support...') print('It is recommended to use GPU if possible...') image_np, mask_np = preprocess_images(image_path, mask_path, max_dim) @@ -17,22 +32,22 @@ def remove_watermark(image_path, mask_path, max_dim, reg_noise, input_depth, lr, num_channels_down = [128] * 5, num_channels_up = [128] * 5, num_channels_skip = [128] * 5 - ).type(DTYPE) + ).type(DTYPE).to(device) - objective = torch.nn.MSELoss().type(DTYPE) + objective = torch.nn.MSELoss().type(DTYPE).to(device) optimizer = optim.Adam(generator.parameters(), lr) - image_var = np_to_torch_array(image_np).type(DTYPE) - mask_var = np_to_torch_array(mask_np).type(DTYPE) + image_var = np_to_torch_array(image_np).type(DTYPE).to(device) + mask_var = np_to_torch_array(mask_np).type(DTYPE).to(device) - generator_input = input_noise(input_depth, image_np.shape[1:]).type(DTYPE) + generator_input = input_noise(input_depth, image_np.shape[1:]).type(DTYPE).to(device) generator_input_saved = generator_input.detach().clone() noise = generator_input.detach().clone() print('\nStarting training...\n') - progress_bar = tqdm(range(training_steps), desc = 'Completed', ncols = tqdm_length) + progress_bar = tqdm(range(training_steps), desc='Completed', ncols=tqdm_length) for step in progress_bar: optimizer.zero_grad() @@ -40,20 +55,20 @@ def remove_watermark(image_path, mask_path, max_dim, reg_noise, input_depth, lr, if reg_noise > 0: generator_input = generator_input_saved + (noise.normal_() * reg_noise) - + output = generator(generator_input) - + loss = objective(output * mask_var, image_var * mask_var) loss.backward() if step % show_step == 0: output_image = torch_to_np_array(output) visualize_sample(image_np, output_image, nrow = 2, size_factor = 10) - + progress_bar.set_postfix(Loss = loss.item()) - + optimizer.step() - + output_image = torch_to_np_array(output) visualize_sample(output_image, nrow = 1, size_factor = 10) @@ -62,4 +77,4 @@ def remove_watermark(image_path, mask_path, max_dim, reg_noise, input_depth, lr, output_path = image_path.split('/')[-1].split('.')[-2] + '-output.jpg' print(f'\nSaving final output image to: "{output_path}"\n') - pil_image.save(output_path) \ No newline at end of file + pil_image.save(output_path)