Repository: pytorch
Issue: Torch Dynamo support for Flux Transformer model
Labels: module: onnx,triaged

### 🐛 Describe the bug

I'm using the script below to export the Flux Transformer model to ONNX using torch.onnx.dynamo_export(). However, I run into a TypeError relating to an attribute type. 

The script below can be used to reproduce the issue:

```
import torch
from diffusers.models import FluxTransformer2DModel

# Load Model
model_dir = "black-forest-labs/FLUX.1-dev"
device = "cuda"
model = FluxTransformer2DModel.from_pretrained(model_dir, subfolder="transformer", torch_dtype=torch.float16).to(device)

# Define Inputs
dtype = torch.float32
tensor_dtype = torch.float16
batch_size = 1
text_maxlen = 512
latent_height, latent_width = 1024 // 8, 1024 // 8
config = {
    'in_channels': 64,
    'joint_attention_dim': 4096,
    'pooled_projection_dim': 768
}
inputs = {
    'hidden_states': torch.randn(batch_size, (latent_height // 2) * (latent_width // 2), config['in_channels'], dtype=tensor_dtype, device=device),
    'encoder_hidden_states': torch.randn(batch_size, text_maxlen, config['joint_attention_dim'], dtype=tensor_dtype, device=device),
    'pooled_projections': torch.randn(batch_size, config['pooled_projection_dim'], dtype=tensor_dtype, device=device),
    'timestep': torch.tensor([1.]*batch_size, dtype=tensor_dtype, device=device),
    'img_ids': torch.randn((latent_height // 2) * (latent_width // 2), 3, dtype=dtype, device=device),
    'txt_ids': torch.randn(text_maxlen, 3, dtype=dtype, device=device),
    'guidance': torch.tensor([1.]*batch_size, dtype=dtype, device=device),
}

# Export to ONNX
out = torch.onnx.dynamo_export(model,
    **inputs,
)
out.save("transformer_dynamo.onnx")
```

The error is pasted below:
```
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/_exporter_legacy.py", line 801, in dynamo_export
    ).export()
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/_exporter_legacy.py", line 570, in export
    onnxscript_graph = fx_interpreter.run(
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/decorator.py", line 146, in wrapper
    ctx.log_and_raise_if_error(diag)
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/context.py", line 355, in log_and_raise_if_error
    raise diagnostic.source_exception
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/decorator.py", line 130, in wrapper
    return_values = fn(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/fx/fx_onnx_interpreter.py", line 537, in run
    self.run_node(
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/decorator.py", line 146, in wrapper
    ctx.log_and_raise_if_error(diag)
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/context.py", line 355, in log_and_raise_if_error
    raise diagnostic.source_exception
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/decorator.py", line 130, in wrapper
    return_values = fn(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/fx/fx_onnx_interpreter.py", line 450, in run_node
    self.call_module(
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/fx/fx_onnx_interpreter.py", line 738, in call_module
    sub_onnxscript_graph = self.run(
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/decorator.py", line 146, in wrapper
    ctx.log_and_raise_if_error(diag)
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/context.py", line 355, in log_and_raise_if_error
    raise diagnostic.source_exception
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/decorator.py", line 130, in wrapper
    return_values = fn(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/fx/fx_onnx_interpreter.py", line 537, in run
    self.run_node(
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/decorator.py", line 146, in wrapper
    ctx.log_and_raise_if_error(diag)
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/context.py", line 355, in log_and_raise_if_error
    raise diagnostic.source_exception
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/decorator.py", line 130, in wrapper
    return_values = fn(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/fx/fx_onnx_interpreter.py", line 450, in run_node
    self.call_module(
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/fx/fx_onnx_interpreter.py", line 738, in call_module
    sub_onnxscript_graph = self.run(
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/decorator.py", line 146, in wrapper
    ctx.log_and_raise_if_error(diag)
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/context.py", line 355, in log_and_raise_if_error
    raise diagnostic.source_exception
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/decorator.py", line 130, in wrapper
    return_values = fn(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/fx/fx_onnx_interpreter.py", line 537, in run
    self.run_node(
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/decorator.py", line 146, in wrapper
    ctx.log_and_raise_if_error(diag)
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/context.py", line 355, in log_and_raise_if_error
    raise diagnostic.source_exception
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/diagnostics/infra/decorator.py", line 130, in wrapper
    return_values = fn(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/fx/fx_onnx_interpreter.py", line 440, in run_node
    self.call_function(
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/fx/fx_onnx_interpreter.py", line 656, in call_function
    ) = symbolic_fn(*onnx_args, **onnx_kwargs)
  File "/usr/local/lib/python3.10/dist-packages/onnxscript/values.py", line 583, in __call__
    return self.func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/onnxscript/function_libs/torch_lib/ops/nn.py", line 1883, in aten__scaled_dot_product_flash_attention
    ) = _aten__scaled_dot_product_flash_attention_fillin_empty_outputs(query)
  File "/usr/local/lib/python3.10/dist-packages/onnxscript/function_libs/torch_lib/ops/nn.py", line 1842, in _aten__scaled_dot_product_flash_attention_fillin_empty_outputs
    op.Constant(value=onnx.helper.make_tensor("Empty_INTS", INT64.dtype, [0], []))
  File "/usr/local/lib/python3.10/dist-packages/onnxscript/onnx_opset/_impl/opset13.py", line 453, in Constant
    return op(
  File "/usr/local/lib/python3.10/dist-packages/onnxscript/values.py", line 301, in __call__
    return evaluator.default().eval(schema, args, kwargs)
  File "/usr/local/lib/python3.10/dist-packages/onnxscript/function_libs/torch_lib/graph_building/_graph_building_torch.py", line 346, in eval
    return self._graph.add_op_call(schema, inputs, attributes)
  File "/usr/local/lib/python3.10/dist-packages/onnxscript/function_libs/torch_lib/graph_building/_graph_building_torch.py", line 888, in add_op_call
    result = self._add_torchscript_op_call(
  File "/usr/local/lib/python3.10/dist-packages/onnxscript/function_libs/torch_lib/graph_building/_graph_building_torch.py", line 768, in _add_torchscript_op_call
    result = _create_op_call_in_torch_graph(
  File "/usr/local/lib/python3.10/dist-packages/onnxscript/function_libs/torch_lib/graph_building/_graph_building_torch.py", line 508, in _create_op_call_in_torch_graph
    _add_attribute_to_torchscript_node(node, key, value)
  File "/usr/local/lib/python3.10/dist-packages/onnxscript/function_libs/torch_lib/graph_building/_graph_building_torch.py", line 470, in _add_attribute_to_torchscript_node
    raise TypeError(
TypeError: Unsupported attribute type '<class 'onnx.onnx_ml_pb2.TensorProto'>' for attribute 'value' in node=%146 : Tensor = onnx::Constant()
, value is dims: 0
data_type: 7
name: "Empty_INTS"
```

### Versions

transformers              4.42.2
torch                     2.6.0.dev20241016+cu124
diffusers                 0.31.0.dev0