Feat (script): Added calibration size argument.

minimal_script.py

@@ -84,6 +84,8 @@ def main(args):
     dtype = getattr(torch, args.dtype)
 
     calibration_prompts = load_calib_prompts(args.calibration_prompt_path)
+    assert args.calibration_prompts <= len(calibration_prompts) , f"--calibration-prompts must be <= {len(calibration_prompts)}"
+    calibration_prompts = calibration_prompts[:args.calibration_prompts]
     latents = torch.load(args.path_to_latents).to(torch.float16)
 
     # Create output dir. Move to tmp if None
@@ -221,6 +223,11 @@ if __name__ == "__main__":
         '-d', '--device', type=str, default='cuda:0', help='Target device for quantized model.')
     parser.add_argument(
         '--calibration-prompt-path', type=str, default=None, help='Path to calibration prompt')
+    parser.add_argument(
+        '--calibration-prompts',
+        type=int,
+        default=500,
+        help='Number of prompts to use for calibration. Default: %(default)s')
     parser.add_argument(
         '--checkpoint-name',
         type=str,
@@ -234,7 +241,6 @@ if __name__ == "__main__":
         default=None,
         help=
         'Load pre-defined latents. If not provided, they are generated based on an internal seed.')
-
     parser.add_argument('--guidance-scale', type=float, default=8., help='Guidance scale.')
     parser.add_argument(
         '--calibration-steps', type=float, default=8, help='Steps used during calibration')