Early exit in case of incompatibility

src/brevitas_examples/stable_diffusion/main.py

@@ -71,6 +71,11 @@ def run_val_inference(pipe, resolution, prompts, seeds, output_path, device, dty
 
 def main(args):
 
+    if args.export_target:
+        assert args.weight_quant_format == 'int', "Currently only integer quantization supported for export."
+        if args.is_sd_xl:
+            assert args.export_target != 'torchscript', "Torchscript export of SD-XL not supported"
+
     # Select dtype
     if args.float16:
         dtype = torch.float16
@@ -216,8 +221,6 @@ def input_bit_width(module):
             pipe, args.resolution, prompts, test_seeds, output_dir, args.device, dtype)
 
     if args.export_target:
-        assert args.weight_quant_format == 'int', "Only integer quantization supported for export."
-
         # Move to cpu and to float32 to enable CPU export
         if not (args.float16 and args.export_cuda_float16):
             pipe.unet.to('cpu').to(torch.float32)
@@ -227,7 +230,6 @@ def input_bit_width(module):
 
         # Define tracing input
         if args.is_sd_xl:
-            assert args.export_target != 'torchscript', "Torchscript export of SD-XL not supported"
             generate_fn = generate_unet_xl_rand_inputs
             shape = SD_XL_EMBEDDINGS_SHAPE
         else: