Function get_clip_model

examples/diffusion/python_flux/models.py:202–270

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

def get_clip_model(local_dir,
                   hf_model_path,
                   compiled_dir,
                   model_dir="text_encoder",
                   torch_dtype=torch.float32,
                   bs=1,
                   exhaustive_tune=False,
                   fp16=False,
                   bf16=True):
    clip_local_dir = get_local_path(local_dir, model_dir)
    onnx_file = "model.onnx"
    onnx_path = os.path.join(clip_local_dir, onnx_file)

    def get_compiled_file_name():
        name = f"model_b{bs}"
        if fp16:
            name += "_fp16"
        elif bf16:
            name += "_bf16"

        if exhaustive_tune: name += "_exh"
        return name + ".mxr"

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

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

    sample_inputs = (torch.zeros(bs, 77, dtype=torch.int32), )
    input_names = ["input_ids"]
    if not os.path.isfile(onnx_path):
        print("ONNX file not found.. exporting CLIP encoder to ONNX")
        model = CLIPTextModel.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'}}

        # CLIP export requires nightly pytorch due to bug in onnx parser
        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)