[LLM Bench] Allow Image Generation Models to Run in BF16 (#1368)

This change allows setting image generation models to BF16 using config passed while running benchmark. Co-authored-by: Ekaterina Aidova <[email protected]> Co-authored-by: guozhong wang <[email protected]>
openvinotoolkit · Dec 23, 2024 · 1179cb6 · 1179cb6
1 parent eac4f37
commit 1179cb6
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tools/llm_bench/llm_bench_utils/pt_utils.py b/tools/llm_bench/llm_bench_utils/pt_utils.py
@@ -131,6 +131,7 @@ def create_image_gen_model(model_path, device, **kwargs):
             model_class = PT_MODEL_CLASSES_MAPPING[model_type]
             start = time.perf_counter()
             pipe = model_class.from_pretrained(model_path)
+            pipe = set_bf16(pipe, device, **kwargs)
             end = time.perf_counter()
             from_pretrain_time = end - start
         else: