Integrate CharacterGen & Improve its workflow with Unique3D

shared_utils/common_utils.py

@@ -52,13 +52,13 @@ def get_list_filenames(directory, extension_filter=None):
         return []
     
 # Download pre-trained model if it not exist locally
-def resume_or_download_model_from_hf(checkpoints_dir_abs, default_repo_id, model_name, class_name=""):
+def resume_or_download_model_from_hf(checkpoints_dir_abs, repo_id, model_name, class_name="", repo_type="model"):
     
     ckpt_path = os.path.join(checkpoints_dir_abs, model_name)
     if not os.path.isfile(ckpt_path):
-        cstr(f"[{class_name}] can't find checkpoint {ckpt_path}, will download it from repo {default_repo_id} instead").warning.print()
+        cstr(f"[{class_name}] can't find checkpoint {ckpt_path}, will download it from repo {repo_id} instead").warning.print()
         
         from huggingface_hub import hf_hub_download
-        hf_hub_download(repo_id=default_repo_id, local_dir=checkpoints_dir_abs, filename=model_name, repo_type="model")
+        hf_hub_download(repo_id=repo_id, local_dir=checkpoints_dir_abs, filename=model_name, repo_type=repo_type)
 
     return ckpt_path

shared_utils/image_utils.py

@@ -13,7 +13,7 @@ def prepare_torch_img(img, size_H, size_W, device="cuda", keep_shape=False):
         img_new = img_new.permute(0, 2, 3, 1)
     return img_new
 
-def torch_imgs_to_pils(images, masks=None):
+def torch_imgs_to_pils(images, masks=None, alpha_min=0.1):
     """
         images (torch): [N, H, W, C] or [H, W, C]
         masks (torch): [N, H, W] or [H, W]
@@ -24,8 +24,11 @@ def torch_imgs_to_pils(images, masks=None):
     if masks is not None:
         if len(masks.shape) == 2:
             masks = masks.unsqueeze(0)
-        masks = masks.unsqueeze(3)
 
+        inv_mask_index = masks < alpha_min
+        images[inv_mask_index] = 0.
+        
+        masks = masks.unsqueeze(3)
         images = torch.cat((images, masks), dim=3)
         mode="RGBA"
     else: