# sudo cog push r8.im/yael-vinker/clipasso # Prediction interface for Cog ⚙️ # https://github.com/replicate/cog/blob/main/docs/python.md import warnings warnings.filterwarnings('ignore') warnings.simplefilter('ignore') from cog import BasePredictor, Input, Path import subprocess as sp import os import re import imageio import matplotlib.pyplot as plt import numpy as np import pydiffvg import torch from PIL import Image import multiprocessing as mp from shutil import copyfile import argparse import math import sys import time import traceback import PIL import torch.nn as nn import torch.nn.functional as F import wandb from torchvision import models, transforms from tqdm import tqdm import config import sketch_utils as utils from models.loss import Loss from models.painter_params import Painter, PainterOptimizer class Predictor(BasePredictor): def setup(self): """Load the model into memory to make running multiple predictions efficient""" self.num_iter = 2001 self.save_interval = 100 self.num_sketches = 3 self.use_gpu = True def predict( self, target_image: Path = Input(description="Input image (square, without background)"), num_strokes: int = Input(description="The number of strokes used to create the sketch, which determines the level of abstraction",default=16), trials: int = Input(description="It is recommended to use 3 trials to recieve the best sketch, but it might be slower",default=3), mask_object: int = Input(description="It is recommended to use images without a background, however, if your image contains a background, you can mask it out by using this flag with 1 as an argument",default=0), fix_scale: int = Input(description="If your image is not squared, it might be cut off, it is recommended to use this flag with 1 as input to automatically fix the scale without cutting the image",default=0), ) -> Path: self.num_sketches = trials target_image_name = os.path.basename(str(target_image)) multiprocess = False abs_path = os.path.abspath(os.getcwd()) target = str(target_image) assert os.path.isfile(target), f"{target} does not exists!" test_name = os.path.splitext(target_image_name)[0] output_dir = f"{abs_path}/output_sketches/{test_name}/" if not os.path.exists(output_dir): os.makedirs(output_dir) print("=" * 50) print(f"Processing [{target_image_name}] ...") print(f"Results will be saved to \n[{output_dir}] ...") print("=" * 50) if not torch.cuda.is_available(): self.use_gpu = False print("CUDA is not configured with GPU, running with CPU instead.") print("Note that this will be very slow, it is recommended to use colab.") print(f"GPU: {self.use_gpu}") seeds = list(range(0, self.num_sketches * 1000, 1000)) losses_all = {} for seed in seeds: wandb_name = f"{test_name}_{num_strokes}strokes_seed{seed}" sp.run(["python", "config.py", target, "--num_paths", str(num_strokes), "--output_dir", output_dir, "--wandb_name", wandb_name, "--num_iter", str(self.num_iter), "--save_interval", str(self.save_interval), "--seed", str(seed), "--use_gpu", str(int(self.use_gpu)), "--fix_scale", str(fix_scale), "--mask_object", str(mask_object), "--mask_object_attention", str( mask_object), "--display_logs", str(int(0))]) config_init = np.load(f"{output_dir}/{wandb_name}/config_init.npy", allow_pickle=True)[()] args = Args(config_init) args.cog_display = True final_config = vars(args) try: configs_to_save = main(args) except BaseException as err: print(f"Unexpected error occurred:\n {err}") print(traceback.format_exc()) sys.exit(1) for k in configs_to_save.keys(): final_config[k] = configs_to_save[k] np.save(f"{args.output_dir}/config.npy", final_config) if args.use_wandb: wandb.finish() config = np.load(f"{output_dir}/{wandb_name}/config.npy", allow_pickle=True)[()] loss_eval = np.array(config['loss_eval']) inds = np.argsort(loss_eval) losses_all[wandb_name] = loss_eval[inds][0] # return Path(f"{output_dir}/{wandb_name}/best_iter.svg") sorted_final = dict(sorted(losses_all.items(), key=lambda item: item[1])) copyfile(f"{output_dir}/{list(sorted_final.keys())[0]}/best_iter.svg", f"{output_dir}/{list(sorted_final.keys())[0]}_best.svg") target_path = f"{abs_path}/target_images/{target_image_name}" svg_files = os.listdir(output_dir) svg_files = [f for f in svg_files if "best.svg" in f] svg_output_path = f"{output_dir}/{svg_files[0]}" sketch_res = read_svg(svg_output_path, multiply=True).cpu().numpy() sketch_res = Image.fromarray((sketch_res * 255).astype('uint8'), 'RGB') sketch_res.save(f"{abs_path}/output_sketches/sketch.png") return Path(svg_output_path) class Args(): def __init__(self, config): for k in config.keys(): setattr(self, k, config[k]) def load_renderer(args, target_im=None, mask=None): renderer = Painter(num_strokes=args.num_paths, args=args, num_segments=args.num_segments, imsize=args.image_scale, device=args.device, target_im=target_im, mask=mask) renderer = renderer.to(args.device) return renderer def get_target(args): target = Image.open(args.target) if target.mode == "RGBA": # Create a white rgba background new_image = Image.new("RGBA", target.size, "WHITE") # Paste the image on the background. new_image.paste(target, (0, 0), target) target = new_image target = target.convert("RGB") masked_im, mask = utils.get_mask_u2net(args, target) if args.mask_object: target = masked_im if args.fix_scale: target = utils.fix_image_scale(target) transforms_ = [] if target.size[0] != target.size[1]: transforms_.append(transforms.Resize( (args.image_scale, args.image_scale), interpolation=PIL.Image.BICUBIC)) else: transforms_.append(transforms.Resize( args.image_scale, interpolation=PIL.Image.BICUBIC)) transforms_.append(transforms.CenterCrop(args.image_scale)) transforms_.append(transforms.ToTensor()) data_transforms = transforms.Compose(transforms_) target_ = data_transforms(target).unsqueeze(0).to(args.device) return target_, mask def main(args): loss_func = Loss(args) inputs, mask = get_target(args) utils.log_input(args.use_wandb, 0, inputs, args.output_dir) renderer = load_renderer(args, inputs, mask) optimizer = PainterOptimizer(args, renderer) counter = 0 configs_to_save = {"loss_eval": []} best_loss, best_fc_loss = 100, 100 best_iter, best_iter_fc = 0, 0 min_delta = 1e-5 terminate = False renderer.set_random_noise(0) img = renderer.init_image(stage=0) optimizer.init_optimizers() for epoch in tqdm(range(args.num_iter)): renderer.set_random_noise(epoch) if args.lr_scheduler: optimizer.update_lr(counter) start = time.time() optimizer.zero_grad_() sketches = renderer.get_image().to(args.device) losses_dict = loss_func(sketches, inputs.detach( ), renderer.get_color_parameters(), renderer, counter, optimizer) loss = sum(list(losses_dict.values())) loss.backward() optimizer.step_() if epoch % args.save_interval == 0: utils.plot_batch(inputs, sketches, f"{args.output_dir}/jpg_logs", counter, use_wandb=args.use_wandb, title=f"iter{epoch}.jpg") renderer.save_svg( f"{args.output_dir}/svg_logs", f"svg_iter{epoch}") # if args.cog_display: # yield Path(f"{args.output_dir}/svg_logs/svg_iter{epoch}.svg") if epoch % args.eval_interval == 0: with torch.no_grad(): losses_dict_eval = loss_func(sketches, inputs, renderer.get_color_parameters( ), renderer.get_points_parans(), counter, optimizer, mode="eval") loss_eval = sum(list(losses_dict_eval.values())) configs_to_save["loss_eval"].append(loss_eval.item()) for k in losses_dict_eval.keys(): if k not in configs_to_save.keys(): configs_to_save[k] = [] configs_to_save[k].append(losses_dict_eval[k].item()) if args.clip_fc_loss_weight: if losses_dict_eval["fc"].item() < best_fc_loss: best_fc_loss = losses_dict_eval["fc"].item( ) / args.clip_fc_loss_weight best_iter_fc = epoch # print( # f"eval iter[{epoch}/{args.num_iter}] loss[{loss.item()}] time[{time.time() - start}]") cur_delta = loss_eval.item() - best_loss if abs(cur_delta) > min_delta: if cur_delta < 0: best_loss = loss_eval.item() best_iter = epoch terminate = False utils.plot_batch( inputs, sketches, args.output_dir, counter, use_wandb=args.use_wandb, title="best_iter.jpg") renderer.save_svg(args.output_dir, "best_iter") if args.use_wandb: wandb.run.summary["best_loss"] = best_loss wandb.run.summary["best_loss_fc"] = best_fc_loss wandb_dict = {"delta": cur_delta, "loss_eval": loss_eval.item()} for k in losses_dict_eval.keys(): wandb_dict[k + "_eval"] = losses_dict_eval[k].item() wandb.log(wandb_dict, step=counter) if abs(cur_delta) <= min_delta: if terminate: break terminate = True if counter == 0 and args.attention_init: utils.plot_atten(renderer.get_attn(), renderer.get_thresh(), inputs, renderer.get_inds(), args.use_wandb, "{}/{}.jpg".format( args.output_dir, "attention_map"), args.saliency_model, args.display_logs) if args.use_wandb: wandb_dict = {"loss": loss.item(), "lr": optimizer.get_lr()} for k in losses_dict.keys(): wandb_dict[k] = losses_dict[k].item() wandb.log(wandb_dict, step=counter) counter += 1 renderer.save_svg(args.output_dir, "final_svg") path_svg = os.path.join(args.output_dir, "best_iter.svg") utils.log_sketch_summary_final( path_svg, args.use_wandb, args.device, best_iter, best_loss, "best total") return configs_to_save def read_svg(path_svg, multiply=False): device = torch.device("cuda" if ( torch.cuda.is_available() and torch.cuda.device_count() > 0) else "cpu") canvas_width, canvas_height, shapes, shape_groups = pydiffvg.svg_to_scene( path_svg) if multiply: canvas_width *= 2 canvas_height *= 2 for path in shapes: path.points *= 2 path.stroke_width *= 2 _render = pydiffvg.RenderFunction.apply scene_args = pydiffvg.RenderFunction.serialize_scene( canvas_width, canvas_height, shapes, shape_groups) img = _render(canvas_width, # width canvas_height, # height 2, # num_samples_x 2, # num_samples_y 0, # seed None, *scene_args) img = img[:, :, 3:4] * img[:, :, :3] + \ torch.ones(img.shape[0], img.shape[1], 3, device=device) * (1 - img[:, :, 3:4]) img = img[:, :, :3] return img