2model

app.py

@@ -121,13 +121,13 @@ state_dict = torch.load(checkpoint_path, map_location=device)
 nnet_1.load_state_dict(state_dict)
 nnet_1.eval()
 
-# filename = "pretrained_models/t2i_512px_clip_dimr.pth"
-# checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)
-# nnet_2 = utils.get_nnet(**config_2.nnet)
-# nnet_2 = nnet_2.to(device)
-# state_dict = torch.load(checkpoint_path, map_location=device)
-# nnet_2.load_state_dict(state_dict)
-# nnet_2.eval()
+filename = "pretrained_models/t2i_512px_clip_dimr.pth"
+checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)
+nnet_2 = utils.get_nnet(**config_2.nnet)
+nnet_2 = nnet_2.to(device)
+state_dict = torch.load(checkpoint_path, map_location=device)
+nnet_2.load_state_dict(state_dict)
+nnet_2.eval()
 
 # Initialize text model.
 llm = "clip"
@@ -181,11 +181,14 @@ def infer(
     else:
         assert num_of_interpolation == 3, "For arithmetic, please sample three images."
 
-    # if num_of_interpolation == 3:
-    #     nnet = nnet_2
-    # else:
-    #     nnet = nnet_1
-    nnet = nnet_1
+    if num_of_interpolation == 3:
+        nnet = nnet_2
+        config = config_2
+    else:
+        nnet = nnet_1
+        config = config_1
+    # nnet = nnet_1
+    # config = config_1
     
     # Get text embeddings and tokens.
     _context, _token_mask, _token, _caption = get_caption(