update demo

.gitignore

@@ -158,3 +158,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+checkpoints/*
+embeddings/*
+imgs/cropped/*

app.py

@@ -5,7 +5,8 @@ import torch
 import math
 from torchvision import transforms
 from run_pti import run_PTI
-device = "cpu"
+from configs import global_config, paths_config
+device = global_config.device
 years = [str(y) for y in range(1880, 2020, 10)]
 decades = [y + "s" for y in years]
 
@@ -32,14 +33,15 @@ def run_alignment(image_path,idx=None):
     return aligned_image 
 
 def predict(inp, in_decade):
+  in_year = in_decade[:-1]
   #with torch.no_grad():     
   inp.save("imgs/input.png")
   inversion = run_alignment("imgs/input.png", idx=0)
   inversion.save("imgs/cropped/input.png")
-  run_PTI(run_name="gradio_demo", use_wandb=False, use_multi_id_training=False)
+  run_PTI(run_name="gradio_demo", in_year=in_year, use_wandb=False, use_multi_id_training=False)
   #inversion = Image.open("imgs/cropped/input.png")
   
-  in_year = in_decade[:-1]
+  
   pti_models = {}
 
   for year in years:
@@ -57,14 +59,16 @@ def predict(inp, in_decade):
                   p += delta 
 
   space = 0
-  dst = Image.new("RGB", (256 * (len(years) + 1) + (space * len(years)), 256), color='white')
+  #dst = Image.new("RGB", (256 * (len(years) + 1) + (space * len(years)), 256), color='white')
 
 
-  w_pti = torch.load(f"embeddings/{in_year}/PTI/input/0.pt", map_location=device)  
+  w_pti = torch.load(f"embeddings/gradio/PTI/input/0.pt", map_location=device)  
 
   border_width = 10
   #fill_color = 'red'
-  dst.paste(inversion, (0, 0))
+  #dst.paste(inversion, (0, 0))
+  dst = []
+  dst.append(inversion)
 
 
 
@@ -76,14 +80,45 @@ def predict(inp, in_decade):
       #     if year == in_year:
       #         img = img.crop((border_width, border_width, 256 - border_width, 256-border_width))
       #         img = PIL.ImageOps.expand(img, border=border_width, fill=fill_color)
-      dst.paste(img, ((256 + space) * (i+1), 0))  
+      #dst.paste(img, ((256 + space) * (i+1), 0))  
+      dst.append(img)
   dst
   return dst
 
 
-gr.Interface(fn=predict, 
-             inputs=[gr.Image(label="Input Image", type="pil"), gr.Dropdown(label="Input Decade", choices=decades, value="2010s")],
-             outputs=gr.Image(label="Decade Transformations", type="pil"),
-             examples=[["imgs/Steven-Yeun.jpg", "2010s"]]
-             
-             ).launch() #.launch(server_name="0.0.0.0", server_port=8098)
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            
+            in_img = gr.Image(label="Input Image", type="pil")
+            in_year = gr.Dropdown(label="Input Decade", choices=decades, value="2010s")
+            submit = gr.Button(value="Submit")
+            examples = gr.Examples(examples=[["imgs/Steven-Yeun.jpg", "2010s"], ["imgs/00061_1920.png", "1920s"]], inputs=[in_img, in_year])
+        with gr.Column() as outs:
+            with gr.Row():
+                cropped = gr.Image(label=f"Cropped Input", type="pil").style(height=256, width=256)
+                out_1880 = gr.Image(label=f"1880", type="pil").style(height=256, width=256)
+                out_1890 = gr.Image(label=f"1890", type="pil").style(height=256, width=256)                
+            with gr.Row():
+                out_1900 = gr.Image(label=f"1900", type="pil").style(height=256, width=256)
+                out_1910 = gr.Image(label=f"1910", type="pil").style(height=256, width=256)
+                out_1920 = gr.Image(label=f"1920", type="pil").style(height=256, width=256)
+            with gr.Row():
+                out_1930 = gr.Image(label=f"1930", type="pil").style(height=256, width=256)
+                out_1940 = gr.Image(label=f"1940", type="pil").style(height=256, width=256)
+                out_1950 = gr.Image(label=f"1950", type="pil").style(height=256, width=256)
+            with gr.Row():
+                out_1960 = gr.Image(label=f"1960", type="pil").style(height=256, width=256)
+                out_1970 = gr.Image(label=f"1970", type="pil").style(height=256, width=256)
+                out_1980 = gr.Image(label=f"1980", type="pil").style(height=256, width=256)
+            with gr.Row():
+                out_1990 = gr.Image(label=f"1990", type="pil").style(height=256, width=256)
+                out_2000 = gr.Image(label=f"2000", type="pil").style(height=256, width=256)
+                out_2010 = gr.Image(label=f"2010", type="pil").style(height=256, width=256)
+
+    outs = [cropped, out_1880, out_1890, out_1900, out_1910, out_1920, out_1930, out_1940, out_1950, out_1960, out_1970, out_1980, out_1990, out_2000, out_2010]
+    submit.click(predict, inputs=[in_img, in_year], outputs=outs)
+    
+
+demo.launch() #server_name="0.0.0.0", server_port=8098)

checkpoints/model_gradio_demo_input.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:65ee5644ec8ab0966a4eb51971995c071f8178db765750d45e32e0ed18a09738
-size 99867041
+oid sha256:0b43a87853f7d90d798348bdce0934a72d03c2127fec6d28d026e6f023063ce9
+size 99869025

configs/global_config.py

@@ -1,5 +1,5 @@
 ## Device
-cuda_visible_devices = "0"
+cuda_visible_devices = "1"
 device = "cpu"
 
 ## Logs

configs/hyperparameters.py

@@ -12,14 +12,14 @@ regulizer_lpips_lambda = 0.1
 regulizer_alpha = 30
 
 ## Loss
-use_mask = True
+use_mask = False
 pt_l2_lambda = 0.7
 pt_lpips_lambda = 1
 color_transfer_lambda = 0  # 1e6
 id_lambda = 1
 
 ## Steps
-LPIPS_value_threshold = 0.01 # 0.06
+LPIPS_value_threshold = 0.03 # 0.06
 max_pti_steps = 350
 first_inv_steps = 450
 max_images_to_invert = 10

configs/paths_config.py

@@ -4,8 +4,6 @@ year = "2010"
 e4e = "./pretrained_models/e4e_ffhq_encode.pt"
 
 
-stylegan2_ada_ffhq = f"pretrained_models/{year}.pkl"
-
 style_clip_pretrained_mappers = ""
 ir_se50 = "pretrained_models/model_ir_se50.pth"
 dlib = "./pretrained_models/align.dat"
@@ -22,7 +20,7 @@ experiments_output_dir = "./output"
 input_data_path = (
     f"imgs/cropped"
 )
-input_data_id = f"{year}"
+input_data_id = "gradio"
 
 
 

flagged/log.csv

@@ -0,0 +1,2 @@
+Input Image,Input Decade,Cropped Input,1880s,1890s,1900s,1910s,1920s,1930s,1940s,1950s,1960s,1970s,1980s,1990s,2000s,2010s,flag,username,timestamp
+/share/phoenix/nfs04/S7/emc348/faces-through-time/flagged/Input Image/tmpqdqusbe1.png,1920s,/share/phoenix/nfs04/S7/emc348/faces-through-time/flagged/Cropped Input/tmpewip1qud.png,/share/phoenix/nfs04/S7/emc348/faces-through-time/flagged/1880s/tmpqihfju9s.png,/share/phoenix/nfs04/S7/emc348/faces-through-time/flagged/1890s/tmphowwncle.png,/share/phoenix/nfs04/S7/emc348/faces-through-time/flagged/1900s/tmpiyqo8vid.png,/share/phoenix/nfs04/S7/emc348/faces-through-time/flagged/1910s/tmplsvkbwmm.png,/share/phoenix/nfs04/S7/emc348/faces-through-time/flagged/1920s/tmpm88oq3m5.png,/share/phoenix/nfs04/S7/emc348/faces-through-time/flagged/1930s/tmp1xg1cf_k.png,/share/phoenix/nfs04/S7/emc348/faces-through-time/flagged/1940s/tmpyxtlostd.png,/share/phoenix/nfs04/S7/emc348/faces-through-time/flagged/1950s/tmp00ekzg14.png,/share/phoenix/nfs04/S7/emc348/faces-through-time/flagged/1960s/tmp0ptji34n.png,/share/phoenix/nfs04/S7/emc348/faces-through-time/flagged/1970s/tmppty6lfdc.png,/share/phoenix/nfs04/S7/emc348/faces-through-time/flagged/1980s/tmpo1znpbho.png,/share/phoenix/nfs04/S7/emc348/faces-through-time/flagged/1990s/tmpkoptsgzw.png,/share/phoenix/nfs04/S7/emc348/faces-through-time/flagged/2000s/tmpn85t4qg9.png,/share/phoenix/nfs04/S7/emc348/faces-through-time/flagged/2010s/tmpxwhqu3bi.png,,,2023-06-07 16:44:52.259412

run_pti.py

@@ -12,7 +12,7 @@ from training.coaches.single_id_coach import SingleIDCoach
 from utils.ImagesDataset import ImagesDataset
 
 
-def run_PTI(run_name="", use_wandb=False, use_multi_id_training=False):
+def run_PTI(run_name="", in_year="2010", use_wandb=False, use_multi_id_training=False):
     os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
     os.environ["CUDA_VISIBLE_DEVICES"] = global_config.cuda_visible_devices
 
@@ -40,9 +40,9 @@ def run_PTI(run_name="", use_wandb=False, use_multi_id_training=False):
     dataloader = DataLoader(dataset, batch_size=1, shuffle=False)
 
     if use_multi_id_training:
-        coach = MultiIDCoach(dataloader, use_wandb)
+        coach = MultiIDCoach(dataloader, in_year, use_wandb)
     else:
-        coach = SingleIDCoach(dataloader, use_wandb)
+        coach = SingleIDCoach(dataloader, in_year, use_wandb)
 
     coach.train()
 
@@ -50,6 +50,6 @@ def run_PTI(run_name="", use_wandb=False, use_multi_id_training=False):
 
 
 if __name__ == "__main__":
-    run_name = f"pti_{paths_config.year}"        
+    run_name = f"pti"        
     print(run_name)
-    run_PTI(run_name=run_name, use_wandb=False, use_multi_id_training=False)
+    run_PTI(run_name=run_name, in_year="2010", use_wandb=False, use_multi_id_training=False)

training/coaches/base_coach.py

@@ -23,7 +23,7 @@ import copy
 
 
 class BaseCoach:
-    def __init__(self, data_loader, use_wandb):
+    def __init__(self, data_loader, in_year, use_wandb):
 
         self.use_wandb = use_wandb
         self.data_loader = data_loader
@@ -56,6 +56,7 @@ class BaseCoach:
             .to(global_config.device)
             .eval()
         )
+        self.in_year = in_year
 
         if hyperparameters.use_mask:
             self.mask = mask.Mask(device=global_config.device)
@@ -69,11 +70,11 @@ class BaseCoach:
     def restart_training(self):
 
         # Initialize networks
-        self.G = load_old_G()
+        self.G = load_old_G(self.in_year)
 
         toogle_grad(self.G, True)
 
-        self.original_G = load_old_G()
+        self.original_G = load_old_G(self.in_year)
 
         self.space_regulizer = Space_Regulizer(self.original_G, self.lpips_loss)
         self.optimizer = self.configure_optimizers()

training/coaches/multi_id_coach.py

@@ -10,8 +10,8 @@ from utils.log_utils import log_images_from_w
 
 
 class MultiIDCoach(BaseCoach):
-    def __init__(self, data_loader, use_wandb):
-        super().__init__(data_loader, use_wandb)
+    def __init__(self, data_loader, in_year, use_wandb):
+        super().__init__(data_loader, in_year, use_wandb)
 
     def train(self):
         self.G.synthesis.train()

training/coaches/single_id_coach.py

@@ -9,8 +9,8 @@ import copy
 
 
 class SingleIDCoach(BaseCoach):
-    def __init__(self, data_loader, use_wandb):
-        super().__init__(data_loader, use_wandb)
+    def __init__(self, data_loader, in_year, use_wandb):
+        super().__init__(data_loader, in_year, use_wandb)
 
     def train(self):
 

utils/models_utils.py

@@ -18,8 +18,8 @@ def load_tuned_G(run_id, type):
     return new_G
 
 
-def load_old_G():
-    with open(paths_config.stylegan2_ada_ffhq, 'rb') as f:
+def load_old_G(in_year):
+    with open(f"pretrained_models/{in_year}.pkl", 'rb') as f:
         old_G = pickle.load(f)['G_ema'].to(global_config.device).eval()
         old_G = old_G.float()
     return old_G