[sharktank] Fix VAE string formatting in tests (#1014)

Add missing format prefix f.

Author: sogartar

sharktank/tests/models/vae/vae_test.py

@@ -49,42 +49,43 @@ def setUp(self):
         hf_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
         hf_hub_download(
             repo_id=hf_model_id,
-            local_dir="{self._temp_dir}",
+            local_dir=f"{self._temp_dir}",
             local_dir_use_symlinks=False,
             revision="main",
             filename="vae/config.json",
         )
         hf_hub_download(
             repo_id=hf_model_id,
-            local_dir="{self._temp_dir}",
+            local_dir=f"{self._temp_dir}",
             local_dir_use_symlinks=False,
             revision="main",
             filename="vae/diffusion_pytorch_model.safetensors",
         )
         hf_hub_download(
             repo_id="amd-shark/sdxl-quant-models",
-            local_dir="{self._temp_dir}",
+            local_dir=f"{self._temp_dir}",
             local_dir_use_symlinks=False,
             revision="main",
             filename="vae/vae.safetensors",
         )
         torch.manual_seed(12345)
         f32_dataset = import_hf_dataset(
-            "{self._temp_dir}/vae/config.json",
-            ["{self._temp_dir}/vae/diffusion_pytorch_model.safetensors"],
+            f"{self._temp_dir}/vae/config.json",
+            [f"{self._temp_dir}/vae/diffusion_pytorch_model.safetensors"],
         )
-        f32_dataset.save("{self._temp_dir}/vae_f32.irpa", io_report_callback=print)
+        f32_dataset.save(f"{self._temp_dir}/vae_f32.irpa", io_report_callback=print)
         f16_dataset = import_hf_dataset(
-            "{self._temp_dir}/vae/config.json", ["{self._temp_dir}/vae/vae.safetensors"]
+            f"{self._temp_dir}/vae/config.json",
+            [f"{self._temp_dir}/vae/vae.safetensors"],
         )
-        f16_dataset.save("{self._temp_dir}/vae_f16.irpa", io_report_callback=print)
+        f16_dataset.save(f"{self._temp_dir}/vae_f16.irpa", io_report_callback=print)
 
     def testCompareF32EagerVsHuggingface(self):
         dtype = getattr(torch, "float32")
         inputs = get_random_inputs(dtype=dtype, device="cpu", bs=1)
-        ref_results = run_torch_vae("{self._temp_dir}", inputs)
+        ref_results = run_torch_vae(f"{self._temp_dir}", inputs)
 
-        ds = Dataset.load("{self._temp_dir}/vae_f32.irpa", file_type="irpa")
+        ds = Dataset.load(f"{self._temp_dir}/vae_f32.irpa", file_type="irpa")
         model = VaeDecoderModel.from_dataset(ds).to(device="cpu")
 
         results = model.forward(inputs)
@@ -95,9 +96,9 @@ def testCompareF32EagerVsHuggingface(self):
     def testCompareF16EagerVsHuggingface(self):
         dtype = getattr(torch, "float32")
         inputs = get_random_inputs(dtype=dtype, device="cpu", bs=1)
-        ref_results = run_torch_vae("{self._temp_dir}", inputs)
+        ref_results = run_torch_vae(f"{self._temp_dir}", inputs)
 
-        ds = Dataset.load("{self._temp_dir}/vae_f16.irpa", file_type="irpa")
+        ds = Dataset.load(f"{self._temp_dir}/vae_f16.irpa", file_type="irpa")
         model = VaeDecoderModel.from_dataset(ds).to(device="cpu")
 
         results = model.forward(inputs.to(torch.float16))
@@ -107,10 +108,10 @@ def testCompareF16EagerVsHuggingface(self):
     def testVaeIreeVsHuggingFace(self):
         dtype = getattr(torch, "float32")
         inputs = get_random_inputs(dtype=dtype, device="cpu", bs=1)
-        ref_results = run_torch_vae("{self._temp_dir}", inputs)
+        ref_results = run_torch_vae(f"{self._temp_dir}", inputs)
 
-        ds_f16 = Dataset.load("{self._temp_dir}/vae_f16.irpa", file_type="irpa")
-        ds_f32 = Dataset.load("{self._temp_dir}/vae_f32.irpa", file_type="irpa")
+        ds_f16 = Dataset.load(f"{self._temp_dir}/vae_f16.irpa", file_type="irpa")
+        ds_f32 = Dataset.load(f"{self._temp_dir}/vae_f32.irpa", file_type="irpa")
 
         model_f16 = VaeDecoderModel.from_dataset(ds_f16).to(device="cpu")
         model_f32 = VaeDecoderModel.from_dataset(ds_f32).to(device="cpu")
@@ -119,8 +120,8 @@ def testVaeIreeVsHuggingFace(self):
         module_f16 = export_vae(model_f16, inputs.to(torch.float16), True)
         module_f32 = export_vae(model_f32, inputs, True)
 
-        module_f16.save_mlir("{self._temp_dir}/vae_f16.mlir")
-        module_f32.save_mlir("{self._temp_dir}/vae_f32.mlir")
+        module_f16.save_mlir(f"{self._temp_dir}/vae_f16.mlir")
+        module_f32.save_mlir(f"{self._temp_dir}/vae_f32.mlir")
         extra_args = [
             "--iree-hal-target-backends=rocm",
             "--iree-hip-target=gfx942",
@@ -137,22 +138,22 @@ def testVaeIreeVsHuggingFace(self):
         ]
 
         iree.compiler.compile_file(
-            "{self._temp_dir}/vae_f16.mlir",
-            output_file="{self._temp_dir}/vae_f16.vmfb",
+            f"{self._temp_dir}/vae_f16.mlir",
+            output_file=f"{self._temp_dir}/vae_f16.vmfb",
             extra_args=extra_args,
         )
         iree.compiler.compile_file(
-            "{self._temp_dir}/vae_f32.mlir",
-            output_file="{self._temp_dir}/vae_f32.vmfb",
+            f"{self._temp_dir}/vae_f32.mlir",
+            output_file=f"{self._temp_dir}/vae_f32.vmfb",
             extra_args=extra_args,
         )
 
         iree_devices = get_iree_devices(driver="hip", device_count=1)
 
         iree_module, iree_vm_context, iree_vm_instance = load_iree_module(
-            module_path="{self._temp_dir}/vae_f16.vmfb",
+            module_path=f"{self._temp_dir}/vae_f16.vmfb",
             devices=iree_devices,
-            parameters_path="{self._temp_dir}/vae_f16.irpa",
+            parameters_path=f"{self._temp_dir}/vae_f16.irpa",
         )
 
         input_args = OrderedDict([("inputs", inputs.to(torch.float16))])
@@ -178,9 +179,9 @@ def testVaeIreeVsHuggingFace(self):
         )
 
         iree_module, iree_vm_context, iree_vm_instance = load_iree_module(
-            module_path="{self._temp_dir}/vae_f32.vmfb",
+            module_path=f"{self._temp_dir}/vae_f32.vmfb",
             devices=iree_devices,
-            parameters_path="{self._temp_dir}/vae_f32.irpa",
+            parameters_path=f"{self._temp_dir}/vae_f32.irpa",
         )
 
         input_args = OrderedDict([("inputs", inputs)])
@@ -209,30 +210,32 @@ def setUp(self):
         hf_model_id = "black-forest-labs/FLUX.1-dev"
         hf_hub_download(
             repo_id=hf_model_id,
-            local_dir="{self._temp_dir}/flux_vae/",
+            local_dir=f"{self._temp_dir}/flux_vae/",
             local_dir_use_symlinks=False,
             revision="main",
             filename="vae/config.json",
         )
         hf_hub_download(
             repo_id=hf_model_id,
-            local_dir="{self._temp_dir}/flux_vae/",
+            local_dir=f"{self._temp_dir}/flux_vae/",
             local_dir_use_symlinks=False,
             revision="main",
             filename="vae/diffusion_pytorch_model.safetensors",
         )
         torch.manual_seed(12345)
         dataset = import_hf_dataset(
-            "{self._temp_dir}/flux_vae/vae/config.json",
-            ["{self._temp_dir}/flux_vae/vae/diffusion_pytorch_model.safetensors"],
+            f"{self._temp_dir}/flux_vae/vae/config.json",
+            [f"{self._temp_dir}/flux_vae/vae/diffusion_pytorch_model.safetensors"],
         )
-        dataset.save("{self._temp_dir}/flux_vae_bf16.irpa", io_report_callback=print)
+        dataset.save(f"{self._temp_dir}/flux_vae_bf16.irpa", io_report_callback=print)
         dataset_f32 = import_hf_dataset(
-            "{self._temp_dir}/flux_vae/vae/config.json",
-            ["{self._temp_dir}/flux_vae/vae/diffusion_pytorch_model.safetensors"],
+            f"{self._temp_dir}/flux_vae/vae/config.json",
+            [f"{self._temp_dir}/flux_vae/vae/diffusion_pytorch_model.safetensors"],
             target_dtype=torch.float32,
         )
-        dataset_f32.save("{self._temp_dir}/flux_vae_f32.irpa", io_report_callback=print)
+        dataset_f32.save(
+            f"{self._temp_dir}/flux_vae_f32.irpa", io_report_callback=print
+        )
 
     def testCompareBF16EagerVsHuggingface(self):
         dtype = torch.bfloat16
@@ -241,7 +244,7 @@ def testCompareBF16EagerVsHuggingface(self):
             "black-forest-labs/FLUX.1-dev", inputs, 1024, 1024, True, dtype
         )
 
-        ds = Dataset.load("{self._temp_dir}/flux_vae_bf16.irpa", file_type="irpa")
+        ds = Dataset.load(f"{self._temp_dir}/flux_vae_bf16.irpa", file_type="irpa")
         model = VaeDecoderModel.from_dataset(ds).to(device="cpu")
 
         results = model.forward(inputs)
@@ -255,7 +258,7 @@ def testCompareF32EagerVsHuggingface(self):
             "black-forest-labs/FLUX.1-dev", inputs, 1024, 1024, True, dtype
         )
 
-        ds = Dataset.load("{self._temp_dir}/flux_vae_f32.irpa", file_type="irpa")
+        ds = Dataset.load(f"{self._temp_dir}/flux_vae_f32.irpa", file_type="irpa")
         model = VaeDecoderModel.from_dataset(ds).to(device="cpu", dtype=dtype)
 
         results = model.forward(inputs)
@@ -270,8 +273,8 @@ def testVaeIreeVsHuggingFace(self):
             "black-forest-labs/FLUX.1-dev", inputs, 1024, 1024, True, torch.float32
         )
 
-        ds = Dataset.load("{self._temp_dir}/flux_vae_bf16.irpa", file_type="irpa")
-        ds_f32 = Dataset.load("{self._temp_dir}/flux_vae_f32.irpa", file_type="irpa")
+        ds = Dataset.load(f"{self._temp_dir}/flux_vae_bf16.irpa", file_type="irpa")
+        ds_f32 = Dataset.load(f"{self._temp_dir}/flux_vae_f32.irpa", file_type="irpa")
 
         model = VaeDecoderModel.from_dataset(ds).to(device="cpu")
         model_f32 = VaeDecoderModel.from_dataset(ds_f32).to(device="cpu")
@@ -280,8 +283,8 @@ def testVaeIreeVsHuggingFace(self):
         module = export_vae(model, inputs.to(dtype=dtype), True)
         module_f32 = export_vae(model_f32, inputs, True)
 
-        module.save_mlir("{self._temp_dir}/flux_vae_bf16.mlir")
-        module_f32.save_mlir("{self._temp_dir}/flux_vae_f32.mlir")
+        module.save_mlir(f"{self._temp_dir}/flux_vae_bf16.mlir")
+        module_f32.save_mlir(f"{self._temp_dir}/flux_vae_f32.mlir")
 
         extra_args = [
             "--iree-hal-target-backends=rocm",
@@ -299,22 +302,22 @@ def testVaeIreeVsHuggingFace(self):
         ]
 
         iree.compiler.compile_file(
-            "{self._temp_dir}/flux_vae_bf16.mlir",
-            output_file="{self._temp_dir}/flux_vae_bf16.vmfb",
+            f"{self._temp_dir}/flux_vae_bf16.mlir",
+            output_file=f"{self._temp_dir}/flux_vae_bf16.vmfb",
             extra_args=extra_args,
         )
         iree.compiler.compile_file(
-            "{self._temp_dir}/flux_vae_f32.mlir",
-            output_file="{self._temp_dir}/flux_vae_f32.vmfb",
+            f"{self._temp_dir}/flux_vae_f32.mlir",
+            output_file=f"{self._temp_dir}/flux_vae_f32.vmfb",
             extra_args=extra_args,
         )
 
         iree_devices = get_iree_devices(driver="hip", device_count=1)
 
         iree_module, iree_vm_context, iree_vm_instance = load_iree_module(
-            module_path="{self._temp_dir}/flux_vae_bf16.vmfb",
+            module_path=f"{self._temp_dir}/flux_vae_bf16.vmfb",
             devices=iree_devices,
-            parameters_path="{self._temp_dir}/flux_vae_bf16.irpa",
+            parameters_path=f"{self._temp_dir}/flux_vae_bf16.irpa",
         )
 
         input_args = OrderedDict([("inputs", inputs.to(dtype=dtype))])
@@ -339,9 +342,9 @@ def testVaeIreeVsHuggingFace(self):
         )
 
         iree_module, iree_vm_context, iree_vm_instance = load_iree_module(
-            module_path="{self._temp_dir}/flux_vae_f32.vmfb",
+            module_path=f"{self._temp_dir}/flux_vae_f32.vmfb",
             devices=iree_devices,
-            parameters_path="{self._temp_dir}/flux_vae_f32.irpa",
+            parameters_path=f"{self._temp_dir}/flux_vae_f32.irpa",
         )
 
         input_args = OrderedDict([("inputs", inputs)])