updates

README.md

@@ -14,6 +14,7 @@ pip install -r requirements.txt
 
 ```python
 python generate_prompt.py --setting centering-rgba-alfie --fg_prompt 'A photo of a cat with a hat'
+python generate_prompt.py --setting centering-rgba-alfie --fg_prompt 'A large, colorful tree made of money, with lots of yellow and white coins hanging from its branches'
 
 ```
 

alfie/generate.py

@@ -94,6 +94,7 @@ def get_pipe(image_size, scheduler, device):
 def base_arg_parser():
     parser = argparse.ArgumentParser()
     parser.add_argument("--centering", type=str, default='True')
+    parser.add_argument("--k", type=int, default=1)
     parser.add_argument("--resample", type=str, default='True')
     parser.add_argument("--scheduler", type=str, default='euler', choices=['euler', 'euler_ancestral'])
     parser.add_argument("--use_neg_prompt", type=str, default='True')

generate_prompt.py

@@ -101,7 +101,9 @@ def main():
             alfie_rgba_image_filename =  Path(args.save_folder / f"{name}-rgba-alfie.png")
             alfie_rgba_image_filename.parent.mkdir(parents=True, exist_ok=True)
             alpha_mask_alfie = torch.tensor(alpha_mask)
-            alpha_mask_alfie = torch.where(alpha_mask_alfie == 1, normalize_masks(heatmaps['ff_heatmap'] + 1 * heatmaps['cross_heatmap_fg']), 0.)
+            alfa_hat = normalize_masks(heatmaps['ff_heatmap'] + 1 * heatmaps['cross_heatmap_fg'])
+            alfa_hat = (alfa_hat + args.k * alfa_hat).clip(0, 1)
+            alpha_mask_alfie = torch.where(alpha_mask_alfie == 1, alfa_hat, 0.)
             save_rgba(image, alpha_mask_alfie, alfie_rgba_image_filename)
 
         elif args.cutout_model == 'vit-matte':