Skip to content

GitLab

Commit d1d45737 authored by raojy's avatar raojy 💬
Browse files

Update run.py

parent 380f7a99
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
run.py
View file @ d1d45737
...@@ -22,7 +22,7 @@ torch.backends.cudnn.benchmark = False ...@@ -22,7 +22,7 @@ torch.backends.cudnn.benchmark = False
# 加载 pipeline # 加载 pipeline
# 注意：如果你的 DCU 版本（如某些较老的型号）对 bfloat16 支持不佳，可以尝试换成 torch.float16 # 注意：如果你的 DCU 版本（如某些较老的型号）对 bfloat16 支持不佳，可以尝试换成 torch.float16
pipe = ErnieImagePipeline.from_pretrained( pipe = ErnieImagePipeline.from_pretrained(
"/public/home/raojy/project/baidu/ERNIE-Image", "baidu/ERNIE-Image",
torch_dtype=torch.bfloat16 torch_dtype=torch.bfloat16
) )
...@@ -62,4 +62,4 @@ for idx, prompt in enumerate(prompt_list): ...@@ -62,4 +62,4 @@ for idx, prompt in enumerate(prompt_list):
revised_prompt = output.revised_prompts revised_prompt = output.revised_prompts
images = output.images images = output.images
images[0].save(f"../tests/hf_output{idx+1}.png") images[0].save(f"../tests/hf_output{idx+1}.png")
print(f"Prompt {idx+1} revised: {revised_prompt}") print(f"Prompt {idx+1} revised: {revised_prompt}")
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment