fixed training script issues

Author: tianweiy

fastcomposer/data.py

@@ -2,6 +2,11 @@
 import torch
 from torchvision.io import read_image, ImageReadMode
 import glob
+import json
+import numpy as np
+import random
+from copy import deepcopy
+
 
 
 def prepare_image_token_idx(image_token_mask, max_num_objects):

fastcomposer/utils.py

@@ -661,6 +661,14 @@ def parse_args(default=False):
         default=1,
     )
 
+    parser.add_argument(
+        "--text_image_linking",
+        type=str,
+        default="postfuse",
+    )
+
+    parser.add_argument("--freeze_postfuse_module", action="store_true")
+
     if default:
         return parser.parse_args([])