cleanup pipeline test artifacts after completion.

nod-ai · Apr 9, 2024 · fa6ba50 · fa6ba50
1 parent 699ba0d
commit fa6ba50
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/models/turbine_models/tests/sdxl_test.py b/models/turbine_models/tests/sdxl_test.py
@@ -581,6 +581,10 @@ def test05_t2i_generate_images(self):
             arguments["seed"],
         )
         print("Image generation complete.")
+        os.remove(os.path.join(arguments["pipeline_dir"], "prompt_encoder.vmfb"))
+        os.remove(os.path.join(arguments["pipeline_dir"], f"{arguments["scheduler_id"]}_unet_{str(arguments["num_inference_steps"])}.vmfb"))
+        os.remove(os.path.join(arguments["pipeline_dir"], "vae_decode.vmfb"))
+        os.remove(os.path.join(arguments["pipeline_dir"], "full_pipeline.vmfb"))
 
 
 if __name__ == "__main__":