[FR]: Add support for Stable Diffusion 3 on DirectML

[thevishalagarwal]

Proposal Summary

Is there any plan to add support for Stable Diffusion 3 for DirectML?

What component(s) does this request affect?

• OliveModels
• OliveSystems
• OliveEvaluator
• Metrics
• Engine
• Passes
• Other

[Jay19751103]

Hi Olive Team

Does SD3 be included olive optimization process. I tried this by sdxl opt script. It cannot fuse any MHA. MHA have different with previous.

[peterer0625]

Is there a plan?

[saddam213]

You can convert SD3, SD3.5 and other transformer models with Olive

You just need to edit the corresponding config_MODEL_NAME.json files with the required input and output format

SD3 config_unet.json Example

    "input_model": {
        "type": "PyTorchModel",
        "config": {
            "model_path": "M:\\BaseModels\\stable-diffusion-3-medium",
            "model_loader": "unet_load",
            "model_script": "models.py",
            "io_config": {
                "input_names": [ "hidden_states", "timestep", "encoder_hidden_states", "pooled_projections"],
                "output_names": [ "out_sample" ],
                "dynamic_axes": {
                    "hidden_states": {"0": "batch", "1": "channels", "2": "height", "3": "width"},
                    "timestep": {"0": "batch"},
                    "encoder_hidden_states": {"0": "batch", "1": "sequence"}, "pooled_projections": {"0": "batch", "1": "size"}
                }
            },
            "dummy_inputs_func": "unet_conversion_inputs"
        }
    }

SD3 models.py Example (I don't know python so close your eyes, lol)

# -----------------------------------------------------------------------------
# TRANSFORMER
# -----------------------------------------------------------------------------

class WrappedSD3Transformer2DModel(SD3Transformer2DModel):
    def forward(
        self,
        hidden_states: torch.FloatTensor, 
        timestep: torch.FloatTensor, 
        encoder_hidden_states: torch.FloatTensor,
        pooled_projections: torch.FloatTensor
    ) -> Union[SD3Transformer2DModel, Tuple]:

        return super().forward(
            hidden_states = hidden_states,
            encoder_hidden_states = encoder_hidden_states,
            pooled_projections = pooled_projections,
            timestep = timestep,
            block_controlnet_hidden_states = None,
            joint_attention_kwargs = None
        )


def unet_inputs(batchsize, torch_dtype, is_conversion_inputs=False):
    inputs = {
        "hidden_states": torch.rand((batchsize, 16, config.unet_sample_size, config.unet_sample_size), dtype=torch_dtype),
        "timestep": torch.rand((batchsize,), dtype=torch_dtype),
        "encoder_hidden_states": torch.rand((batchsize, 77, 4096), dtype=torch_dtype),
        "pooled_projections": torch.rand((1, 2048), dtype=torch_dtype)
    }
    return inputs


def unet_load(model_name):
    model = WrappedSD3Transformer2DModel.from_pretrained(model_name, subfolder="transformer")
    return model


def unet_conversion_inputs(model=None):
    return tuple(unet_inputs(1, torch.float32, True).values())


def unet_data_loader(data_dir, batchsize, *args, **kwargs):
    return RandomDataLoader(unet_inputs, batchsize, torch.float16)

And just do the same thing for the other models, I only use DirectML but CUDA should also work fine

I do have some public ONNX repos of these conversions on my huggingface
https://huggingface.co/TensorStack/stable-diffusion-3-lite-onnx
https://huggingface.co/TensorStack/stable-diffusion-3.5-large-onnx