cleanup pipeline test artifacts after completion.

nod-ai · Apr 9, 2024 · 0011328 · 0011328
1 parent 699ba0d
commit 0011328
Showing 1 changed file with 12 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,18 @@ 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"],
+                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__":