Update app.py

app.py

@@ -14,6 +14,8 @@ from data import CustomDataLoader
 from data.super_dataset import SuperDataset
 from configs import parse_config
 from utils.augmentation import ImagePathToImage
+import clip
+from torchvision.transforms import Compose, Resize, ToTensor, Normalize, InterpolationMode
 
 
 class Stylizer(nn.Module):
@@ -118,6 +120,8 @@ def tensor2file(input_image):
     else:
         return image_pil
 
+
+device = "cuda"
 def generate_multi_model(input_img):
 
     # parse config
@@ -146,16 +150,15 @@ def generate_multi_model(input_img):
     dataset = SuperDataset(config)
     dataloader = CustomDataLoader(config, dataset)
 
-    device = "cuda"
     model_dict = torch.load("./pretrained_models/phase2_pretrain_90000.pth", map_location='cpu')
 
     # init netG
-    netG = Stylizer(ngf=config['model']['ngf'], phase=2, model_weights=model_dict['G_ema_model']).to(device)
+    model = Stylizer(ngf=config['model']['ngf'], phase=2, model_weights=model_dict['G_ema_model']).to(device)
 
     for data in dataloader:
 
         real_A = data['test_A'].to(device)
-        fake_B = netG(real_A, mixing=False)
+        fake_B = model(real_A, mixing=False)
         output_img = tensor2file(fake_B)  # get image results
     
         return output_img
@@ -167,6 +170,44 @@ def generate_one_shot(src_img, img_prompt):
     output_img = src_img
     return output_img
 
+    # init model
+    state_dict = torch.load(f"./checkpoints/{img_prompt[-2:]}/epoch_latest.pth", map_location='cpu')
+    model = Stylizer(ngf=64, phase=3, model_weights=state_dict['G_ema_model'])
+    model.to(device)
+    model.eval()
+    model.requires_grad_(False)
+
+    clip_model, img_preprocess = clip.load('ViT-B/32', device=args.device)
+    clip_model.eval()
+    clip_model.requires_grad_(False)
+
+    # image transform for stylizer
+    img_transform = Compose([
+        Resize((512, 512), interpolation=InterpolationMode.LANCZOS),
+        ToTensor(),
+        Normalize([0.5], [0.5])
+    ])
+
+    # get clip features
+    with torch.no_grad():
+        img = img_preprocess(Image.open(f"./example/reference/{img_prompt[-2:]}.png")).unsqueeze(0).to(args.device)
+        clip_feats = clip_model.encode_image(img)
+        clip_feats /= clip_feats.norm(dim=1, keepdim=True)
+
+
+    # load image & to tensor
+    img = Image.open(src_img)
+    if not img.mode == 'RGB':
+        img = img.convert('RGB')
+    img = img_transform(img).unsqueeze(0).to(device)
+
+    # stylize it !
+    with torch.no_grad():
+        res = model(img, clip_feats=clip_feats)
+
+    output_img = tensor2file(res)  # get image results
+    return output_img
+
 def generate_zero_shot(src_img, txt_prompt):
     output_img = src_img
     return output_img