Skip to content

Commit

Permalink
fix shapes for timestep
Browse files Browse the repository at this point in the history
  • Loading branch information
eaidova committed Dec 23, 2024
1 parent ca5103d commit 072e1d9
Showing 1 changed file with 6 additions and 0 deletions.
6 changes: 6 additions & 0 deletions optimum/exporters/openvino/model_configs.py
Original file line number Diff line number Diff line change
Expand Up @@ -1812,6 +1812,12 @@ class UnetOpenVINOConfig(UNetOnnxConfig):
DummyUnetTimestepInputGenerator,
) + UNetOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES[2:]

@property
def inputs(self) -> Dict[str, Dict[int, str]]:
common_inputs = super().inputs
common_inputs["timestep"] = {0: "batch_size"}
return common_inputs


@register_in_tasks_manager("sd3-transformer", *["semantic-segmentation"], library_name="diffusers")
@register_in_tasks_manager("sd3-transformer-2d", *["semantic-segmentation"], library_name="diffusers")
Expand Down

0 comments on commit 072e1d9

Please sign in to comment.