set fp16 for inference

cwdhf · Apr 22, 2024 · c053929 · c053929
1 parent c6555b6
commit c053929
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 14 deletions.
diff --git a/inference.py b/inference.py
@@ -220,13 +220,13 @@ def main():
         if args.output_dir is not None:
             os.makedirs(args.output_dir, exist_ok=True)
 
-    weight_dtype = torch.float32
-    if accelerator.mixed_precision == "fp16":
-        weight_dtype = torch.float16
-        args.mixed_precision = accelerator.mixed_precision
-    elif accelerator.mixed_precision == "bf16":
-        weight_dtype = torch.bfloat16
-        args.mixed_precision = accelerator.mixed_precision
+    weight_dtype = torch.float16
+    # if accelerator.mixed_precision == "fp16":
+    #     weight_dtype = torch.float16
+    #     args.mixed_precision = accelerator.mixed_precision
+    # elif accelerator.mixed_precision == "bf16":
+    #     weight_dtype = torch.bfloat16
+    #     args.mixed_precision = accelerator.mixed_precision
 
     # Load scheduler, tokenizer and models.
     noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")

diff --git a/inference_dc.py b/inference_dc.py
@@ -243,13 +243,13 @@ def main():
         if args.output_dir is not None:
             os.makedirs(args.output_dir, exist_ok=True)
 
-    weight_dtype = torch.float32
-    if accelerator.mixed_precision == "fp16":
-        weight_dtype = torch.float16
-        args.mixed_precision = accelerator.mixed_precision
-    elif accelerator.mixed_precision == "bf16":
-        weight_dtype = torch.bfloat16
-        args.mixed_precision = accelerator.mixed_precision
+    weight_dtype = torch.float16
+    # if accelerator.mixed_precision == "fp16":
+    #     weight_dtype = torch.float16
+    #     args.mixed_precision = accelerator.mixed_precision
+    # elif accelerator.mixed_precision == "bf16":
+    #     weight_dtype = torch.bfloat16
+    #     args.mixed_precision = accelerator.mixed_precision
 
     # Load scheduler, tokenizer and models.
     noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")