optimum-cli diffusion policy model issue

[kraza8]

System Info

Hi,
Trying to export a diffusion policy model to onnx format. From the error message and printed list of model types, it looks like “diffusion” model cannot be exported to onnx.
Is there a way to get around this?

optimum-cli export onnx --model lerobot/diffusion_pusht --task reinforcement-learning /onnx/

Traceback (most recent call last):
File "/optimum-cli", line 8, in
sys.exit(main())
File "/python3.10/site-packages/optimum/commands/optimum_cli.py", line 208, in main
service.run()
File "/python3.10/site-packages/optimum/commands/export/onnx.py", line 265, in run
main_export(
File "/python3.10/site-packages/optimum/exporters/onnx/main.py", line 272, in main_export
config = AutoConfig.from_pretrained(
File "/python3.10/site-packages/transformers/models/auto/configuration_auto.py", line 1008, in from_pretrained
raise ValueError(
ValueError: Unrecognized model in lerobot/diffusion_pusht. Should have a model_type key in its config.json, or contain one of the following strings in its name:

Model type form config.json:
"type": "diffusion"

Supported Models:
albert, align, altclip, audio-spectrogram-transformer, autoformer, bark, bart, beit, bert, bert-generation, big_bird, bigbird_pegasus, biogpt, bit, blenderbot, blenderbot-small, blip, blip-2, bloom, bridgetower, bros, camembert, canine, chameleon, chinese_clip, chinese_clip_vision_model, clap, clip, clip_vision_model, clipseg, clvp, code_llama, codegen, cohere, conditional_detr, convbert, convnext, convnextv2, cpmant, ctrl, cvt, data2vec-audio, data2vec-text, data2vec-vision, dbrx, deberta, deberta-v2, decision_transformer, deformable_detr, deit, depth_anything, deta, detr, dinat, dinov2, distilbert, donut-swin, dpr, dpt, efficientformer, efficientnet, electra, encodec, encoder-decoder, ernie, ernie_m, esm, falcon, fastspeech2_conformer, flaubert, flava, fnet, focalnet, fsmt, funnel, fuyu, gemma, gemma2, git, glpn, gpt-sw3, gpt2, gpt_bigcode, gpt_neo, gpt_neox, gpt_neox_japanese, gptj, gptsan-japanese, graphormer, grounding-dino, groupvit, hiera, hubert, ibert, idefics, idefics2, imagegpt, informer, instructblip, instructblipvideo, jamba, jetmoe, jukebox, kosmos-2, layoutlm, layoutlmv2, layoutlmv3, led, levit, lilt, llama, llava, llava-next-video, llava_next, longformer, longt5, luke, lxmert, m2m_100, mamba, mamba2, marian, markuplm, mask2former, maskformer, maskformer-swin, mbart, mctct, mega, megatron-bert, mgp-str, mistral, mixtral, mobilebert, mobilenet_v1, mobilenet_v2, mobilevit, mobilevitv2, mpnet, mpt, mra, mt5, musicgen, musicgen_melody, mvp, nat, nemotron, nezha, nllb-moe, nougat, nystromformer, olmo, oneformer, open-llama, openai-gpt, opt, owlv2, owlvit, paligemma, patchtsmixer, patchtst, pegasus, pegasus_x, perceiver, persimmon, phi, phi3, pix2struct, plbart, poolformer, pop2piano, prophetnet, pvt, pvt_v2, qdqbert, qwen2, qwen2_moe, rag, realm, recurrent_gemma, reformer, regnet, rembert, resnet, retribert, roberta, roberta-prelayernorm, roc_bert, roformer, rt_detr, rt_detr_resnet, rwkv, sam, seamless_m4t, seamless_m4t_v2, segformer, seggpt, sew, sew-d, siglip, siglip_vision_model, speech-encoder-decoder, speech_to_text, speech_to_text_2, speecht5, splinter, squeezebert, stablelm, starcoder2, superpoint, swiftformer, swin, swin2sr, swinv2, switch_transformers, t5, table-transformer, tapas, time_series_transformer, timesformer, timm_backbone, trajectory_transformer, transfo-xl, trocr, tvlt, tvp, udop, umt5, unispeech, unispeech-sat, univnet, upernet, van, video_llava, videomae, vilt, vipllava, vision-encoder-decoder, vision-text-dual-encoder, visual_bert, vit, vit_hybrid, vit_mae, vit_msn, vitdet, vitmatte, vits, vivit, wav2vec2, wav2vec2-bert, wav2vec2-conformer, wavlm, whisper, xclip, xglm, xlm, xlm-prophetnet, xlm-roberta, xlm-roberta-xl, xlnet, xmod, yolos, yoso, zoedepth

Thanks


To reproduce
Download model from HF
Use optimum-cli to export the model

Platform
Linux

OS Version
Ubuntu 22.04.4 LTS

ONNX Runtime Installation
Released Package

ONNX Runtime Version or Commit ID
1.21.0

ONNX Runtime API
Python

Architecture
ARM64

Execution Provider
CUDA

Execution Provider Library Version
12.4

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction (minimal, reproducible, runnable)

To reproduce
Download model from HF
Use optimum-cli to export the model

Expected behavior

onnx export to succeed

[IlyasMoutawwakil]

Hi, we never added support for LeRobot library (only transformers, diffusers, timm, senetence-transformers are supported). We can add support f it's something many users are asking for, or if a user is up to open PR adding it.