Function get_t5_model

examples/diffusion/python_flux/models.py:273–339

(local_dir,
                 hf_model_path,
                 compiled_dir,
                 max_len=512,
                 model_dir="text_encoder_2",
                 torch_dtype=torch.float32,
                 bs=1,
                 exhaustive_tune=False,
                 fp16=False,
                 bf16=True)

def get_t5_model(local_dir,
                 hf_model_path,
                 compiled_dir,
                 max_len=512,
                 model_dir="text_encoder_2",
                 torch_dtype=torch.float32,
                 bs=1,
                 exhaustive_tune=False,
                 fp16=False,
                 bf16=True):
    t5_local_dir = get_local_path(local_dir, model_dir)
    onnx_file = "model.onnx"
    onnx_path = os.path.join(t5_local_dir, onnx_file)

    def get_compiled_file_name():
        name = f"model_b{bs}"
        name += f"_l{max_len}"
        if fp16:
            name += "_fp16"
        elif bf16:
            name += "_bf16"
        if exhaustive_tune: name += "_exh"
        return name + ".mxr"

    t5_compiled_dir = get_local_path(compiled_dir, model_dir)
    mxr_file = get_compiled_file_name()
    mxr_path = os.path.join(t5_compiled_dir, mxr_file)

    if os.path.isfile(mxr_path):
        print(f"found compiled model.. loading T5 encoder from {mxr_path}")
        model = MGXModel(mxr_path)
        return model

    sample_inputs = (torch.zeros(bs, max_len, dtype=torch.int32), )
    input_names = ["input_ids"]
    if not os.path.isfile(onnx_path):
        model = T5EncoderModel.from_pretrained(hf_model_path,
                                               subfolder=model_dir,
                                               torch_dtype=torch_dtype)
        output_names = ["text_embeddings"]
        dynamic_axes = {"input_ids": {0: 'B'}, "text_embeddings": {0: 'B'}}

        with torch.inference_mode():
            torch.onnx.export(model,
                              sample_inputs,
                              onnx_path,
                              export_params=True,
                              input_names=input_names,
                              output_names=output_names,
                              dynamic_axes=dynamic_axes)

    assert os.path.isfile(onnx_path)
    print(f"Generating MXR from ONNX file: {onnx_path}")
    input_shapes = {
        n: list(t.size())
        for n, t in zip(input_names, sample_inputs)
    }
    model = MGXModel(onnx_path,