Hidream Model loading takes too long — any way to speed it up?

[Me-verner]

Hi, thanks for this great project.

I'm running Hidream with this library in a serverless environment and facing major delays during model loading. It can be very frustrating, especially for time-sensitive or ephemeral deployments.

I've tried everything I could think of to reduce the loading time, but nothing has worked so far. Does anyone have any tips, tricks, or even sample code to help speed up the model initialization?

Any guidance would be greatly appreciated!

[DN6]

Hi @Me-verner are you using any mounted volumes to load the model checkpoint?

[Me-verner]

Hi @Me-verner are you using any mounted volumes to load the model checkpoint?

Nope, its directly on the local fadt nvme storage

[DN6]

@Me-verner Could you share a code example of how you're trying to load the model? And could you also give me an estimate of the loading time?

[Me-verner]

@Me-verner Could you share a code example of how you're trying to load the model? And could you also give me an estimate of the loading time?

Hi @DN6 , thanks for following up.

Sure, here’s a simplified version of how I’m currently loading the model components in parallel using ThreadPoolExecutor, then building the pipeline. Everything is loaded from local fast NVMe storage using local_files_only=True.

from concurrent.futures import ThreadPoolExecutor
from transformers import PreTrainedTokenizerFast, LlamaForCausalLM
from hi_diffusers import HiDreamImageTransformer2DModel, HiDreamImagePipeline
from hi_diffusers.schedulers.flash_flow_match import FlashFlowMatchEulerDiscreteScheduler
import torch, time

t0 = time.time()
with ThreadPoolExecutor() as ex:
    tokenizer_f = ex.submit(
        PreTrainedTokenizerFast.from_pretrained,
        "meta-llama/Meta-Llama-3.1-8B-Instruct",
        local_files_only=True
    )
    encoder_f = ex.submit(
        LlamaForCausalLM.from_pretrained,
        "meta-llama/Meta-Llama-3.1-8B-Instruct",
        torch_dtype=torch.bfloat16,
        use_safetensors=True,
        low_cpu_mem_usage=True,
        local_files_only=True
    )
    transformer_f = ex.submit(
        HiDreamImageTransformer2DModel.from_pretrained,
        "HiDream-ai/HiDream-I1-Fast",
        subfolder="transformer",
        torch_dtype=torch.bfloat16,
        use_safetensors=True,
        local_files_only=True
    )

    tokenizer = tokenizer_f.result()
    encoder = encoder_f.result().to("cuda")
    transformer = transformer_f.result().to("cuda")

scheduler = FlashFlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=3.0, use_dynamic_shifting=False)
pipe = HiDreamImagePipeline.from_pretrained(
    "HiDream-ai/HiDream-I1-Fast",
    tokenizer_4=tokenizer,
    text_encoder_4=encoder,
    scheduler=scheduler,
    torch_dtype=torch.bfloat16,
    use_safetensors=True,
    local_files_only=True
).to("cuda")
pipe.transformer = transformer
t1 = time.time()

print(f"Total model load + pipeline build time: {t1 - t0:.2f}s")

• Typical timings I’m seeing: Load: 15.66s | Build: 16.40s | Warm: 2.55s | Infer: 9.90s | Total: 44.52s

• The main issue is that I'm in a serverless environment, and every time a Docker container spins up, I need to reload all models from scratch. Keeping a worker warm is costly, so I can't rely on persistent memory or cached state. Any ideas for minimizing cold start times would be super helpful

[JoeGaffney]

Could you pre quantize and cache the transformer and some of the larger text_encoders. I've found with other models it can speed the loading from disk part. But you may still want the 16bit depending on your use case.

But is tricky issue how to handle cold starts well with this much data.

[JoeGaffney]

Possibly autoscaling workers that keep models loaded in memory can reduce both cold start time and costs compared to serverless functions that reload the model each time. Once a worker is warm, it can serve multiple requests quickly, spreading the load overhead and lowering overall compute usage.

I’m not exactly sure what the big inference providers do, but I imagine it’s something along these lines. Obviously, this can be trickier to set up on the infrastructure side to get right.