More fixes and checks

Xilinx · Mar 16, 2024 · 845abb1 · 845abb1
1 parent 043d02b
commit 845abb1
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/src/brevitas_examples/stable_diffusion/main.py b/src/brevitas_examples/stable_diffusion/main.py
@@ -161,11 +161,14 @@ def input_bit_width(module):
     print(f"Moving model to {args.device}...")
     pipe = pipe.to(args.device)
 
-    with calibration_mode(pipe.unet):
-        prompts = VALIDATION_PROMPTS
-        run_val_inference(
-            pipe, args.resolution, prompts, test_seeds, output_dir, args.device, dtype)
-
+    if is_input_quantized and args.input_scale_type == 'static':
+        print("Applying activation calibration")
+        with calibration_mode(pipe.unet):
+            prompts = VALIDATION_PROMPTS
+            run_val_inference(
+                pipe, args.resolution, prompts, test_seeds, output_dir, args.device, dtype)
+
+    print("Applying bias correction")
     with bias_correction_mode(pipe.unet):
         prompts = VALIDATION_PROMPTS
         run_val_inference(