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

Update run.py

parent f0177276
...@@ -5,28 +5,25 @@ import torch ...@@ -5,28 +5,25 @@ import torch
from diffusers import ErnieImagePipeline from diffusers import ErnieImagePipeline
from tqdm import tqdm from tqdm import tqdm
# 设置全局随机种子确保可复现性
# seed = 42
seed = random.randint(0, 100000) seed = random.randint(0, 100000)
print(f"seed: {seed}") print(f"seed: {seed}")
random.seed(seed) random.seed(seed)
np.random.seed(seed) np.random.seed(seed)
torch.manual_seed(seed) torch.manual_seed(seed)
# 在 DCU 上,torch.cuda.manual_seed_all 会自动映射到底层 hipRAND
torch.cuda.manual_seed_all(seed) torch.cuda.manual_seed_all(seed)
# 允许一定的算子融合和自动寻优,DCU 的 MIOpen 会接管 cudnn 的设置
torch.backends.cudnn.deterministic = True torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False torch.backends.cudnn.benchmark = False
# 加载 pipeline
# 注意:如果你的 DCU 版本(如某些较老的型号)对 bfloat16 支持不佳,可以尝试换成 torch.float16
pipe = ErnieImagePipeline.from_pretrained( pipe = ErnieImagePipeline.from_pretrained(
"baidu/ERNIE-Image", "baidu/ERNIE-Image",
torch_dtype=torch.bfloat16 torch_dtype=torch.bfloat16
) )
# DCU 版本的 PyTorch 会自动将 "cuda" 映射到 DCU 设备上
pipe = pipe.to("cuda") pipe = pipe.to("cuda")
pipe.transformer.eval() pipe.transformer.eval()
...@@ -34,13 +31,10 @@ pipe.vae.eval() ...@@ -34,13 +31,10 @@ pipe.vae.eval()
pipe.text_encoder.eval() pipe.text_encoder.eval()
pipe.pe.eval() pipe.pe.eval()
# 如果显存不够可以开启 offload
# pipe.enable_model_cpu_offload()
# 设置随机种子,"cuda" 在这里同样会被映射到 DCU
generator = torch.Generator(device="cuda").manual_seed(seed) generator = torch.Generator(device="cuda").manual_seed(seed)
# 确保输出目录存在
os.makedirs("../tests", exist_ok=True) os.makedirs("../tests", exist_ok=True)
# 生成图片 # 生成图片
......
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