rotary embedding refactor 2: update comments, fix dtype for use_real=False

[yiyixuxu]

also made a slow tests for all pipelines using rotary embeddings

slow test for pipelines using rotary embedding

fromdiffusersimportDiffusionPipelinefrombenchmarks.utilsimportbenchmark_fn, flush, bytes_to_giga_bytes, BenchmarkInfoimportargparseimporttorchimportosfromtypingimportDict, UnionimportcsvBENCHMARK_FIELDS= [
"pipeline_cls",
"repo_id",
"time (secs)",
"memory (gbs)",
]
defwrite_to_csv(file_name: str, data_dict: Dict[str, Union[str, bool, float]]):
"""Append a dictionary into a CSV file."""file_exists=os.path.isfile(file_name)
withopen(file_name, mode="a", newline="") ascsvfile:
writer=csv.DictWriter(csvfile, fieldnames=BENCHMARK_FIELDS)
ifnotfile_exists:
writer.writeheader()
writer.writerow(data_dict)
REPO_ID_MAPPING= {
"flux-schnell": "black-forest-labs/FLUX.1-schnell",
"cogvideox-5b": "THUDM/CogVideoX-5b",
"stable-audio": "stabilityai/stable-audio-open-1.0",
"hunyuan-dit": "Tencent-Hunyuan/HunyuanDiT-Diffusers",
"lumina": "Alpha-VLLM/Lumina-Next-SFT-diffusers",
}
DTYPE_MAPPING= {
"flux-schnell": torch.bfloat16,
"cogvideox-5b": torch.bfloat16,
"stable-audio": torch.float16,
"hunyuan-dit": torch.float16,
"lumina": torch.bfloat16,
}
OUTPUT_TYPE_MAPPING= {
"flux-schnell": "image",
"cogvideox-5b": "video",
"stable-audio": "audio",
"hunyuan-dit": "image",
"lumina": "image",
}
CALL_ARGS_MAPPING= {
"flux-schnell": {
"prompt": "A cat holding a sign that says hello world",
"guidance_scale": 0.0,
"num_inference_steps":4,
"max_sequence_length": 256,
},
"cogvideox-5b": {
"prompt": "A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance.",
"num_inference_steps": 50,
"num_frames": 49,
"guidance_scale": 6,
},
"stable-audio": {
"prompt": "The sound of a hammer hitting a wooden surface.",
"negative_prompt": "Low quality.",
"num_inference_steps": 200,
"audio_end_in_s": 10.0,
"num_waveforms_per_prompt": 3, },
"hunyuan-dit": {
"prompt": "一个宇航员在骑马", },
"lumina": {
"prompt": "Upper body of a young woman in a Victorian-era outfit with brass goggles and leather straps. Background shows an industrial revolution cityscape with smoky skies and tall, metal structures",
},
}
defload_pipeline(model_name, args):
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
pipeline=DiffusionPipeline.from_pretrained(REPO_ID_MAPPING[model_name], torch_dtype=DTYPE_MAPPING[model_name])
ifargs.offload:
pipeline.enable_model_cpu_offload()
else:
pipeline=pipeline.to("cuda")
ifargs.run_compile:
pipeline.transformer.to(memory_format=torch.channels_last)
pipeline.vae.to(memory_format=torch.channels_last)
pipeline.transformer=torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)
pipeline.vae.decode=torch.compile(pipeline.vae.decode, mode="max-autotune", fullgraph=True)
pipeline.set_progress_bar_config(disable=True)
returnpipelineif__name__=="__main__":
parser=argparse.ArgumentParser()
parser.add_argument("--branch", default="main", type=str)
parser.add_argument("--run_compile", action="store_true", help="test with torch.compile")
parser.add_argument("--offload", action="store_true", help="test with enable_model_cpu_offload")
parser.add_argument("--result_dir", default="slow_test_rope", type=str)
args=parser.parse_args()
ifargs.run_compileandargs.offload:
raiseValueError("cannot use both --compile and --offload flags")
ifnotos.path.exists(args.result_dir):
os.makedirs(args.result_dir)
log_filename_prefix=f"{args.result_dir}/{args.branch}_compile@{args.run_compile}_offload@{args.offload}"formodel_name, ckpt_idinREPO_ID_MAPPING.items():
pipeline=load_pipeline(model_name, args)
call_kwargs=CALL_ARGS_MAPPING[model_name]
defrun_inference(pipeline, call_kwargs):
_=pipeline(
**call_kwargs,
generator=torch.Generator("cpu").manual_seed(0),
)
flush()
print(f"[INFO] {model_name}: Running benchmark with: {vars(args)}\n")
time=benchmark_fn(run_inference, pipeline, call_kwargs) # in seconds.memory=bytes_to_giga_bytes(torch.cuda.max_memory_allocated()) # in GBs.benchmark_info=BenchmarkInfo(time=time, memory=memory)
flush()
data_dict= {
"pipeline_cls": model_name, "repo_id": ckpt_id, "time (secs)": benchmark_info.time,
"memory (gbs)": benchmark_info.memory,
}
write_to_csv(log_filename_prefix+".csv", data_dict)
print(f"Log written to {log_filename_prefix+'.cvs'}")
print(f"Memory: {benchmark_info.memory} gbs")
print(f"Execution time: {benchmark_info.time} sec")
out=pipeline(**call_kwargs, generator=torch.Generator("cpu").manual_seed(0), return_dict=False)[0][0]
output_type=OUTPUT_TYPE_MAPPING[model_name]
ifoutput_type=="audio":
importsoundfileassfout=out.T.float().cpu().numpy()
out_file_name=log_filename_prefix+f"_{model_name}.wav"sf.write(out_file_name, out, pipeline.vae.sampling_rate)
elifoutput_type=="video":
fromdiffusers.utilsimportexport_to_videoout_file_name=log_filename_prefix+f"_{model_name}.mp4"export_to_video(out, out_file_name, fps=8)
else:
out_file_name=log_filename_prefix+f"_{model_name}.png"out.save(out_file_name)
print(f" output saved to {out_file_name}")

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.