[DRAFT] AMD/ROCM Support

[Fmstrat]

This DRAFT PR is under development to facilitate conversation about a formalized ROCM implementaiton in Krita AI Diffusion.

Outside of KAD, I have a functioning Dockerfile that supports KAD with all features of SD1.5, SDXL, and FLUX (including Kontext, but not SDVQ quant since Nunchaku does not support ROCM). This PR will port my work directly into KAD so that it can formally support ROCM going forward, but will need some guidance to complete as I'm not familiar with the code base.

Working:

• Full ROCM implementation with external ComfyUI (Tested on MI300X)

In progress:

• Dockerfile based on successful external install
• Checks on supported nodes for models to throw errors on lack of SDVQ (if needed, not yet in PR)

To do:

• UI updates
• Model Database recommendations for non-SDVQ FLUX

I'm going to assume there are a number of things in the to do list that I'm missing. My biggest question atm is there is no documentation on how to use the docker.py script in a development environment to build the container, as the ComfyUI setup is done outside of the Dockerfile (my custom setup does the full build in the Dockerfile).

[Acly]

Having built in ROCM support would be great!

My biggest question atm is there is no documentation on how to use the docker.py script in a development environment to build the container

What is unclear exactly? You run python scripts/docker.py and it will download and set everything up for the docker build.

Regarding Nunchaku/SVDQ:

• docker.py / CUDA Dockerfile is still missing support for Nunchaku anyway, just something to keep in mind
• you can pass --no-cuda to start.sh (which passes it on to download_models.py), this will skip SVDQ models
• might still need non-SVDQ alternatives in model database

For local testing of both the docker image and the UI installer it can be helpful to run the local file server instead of downloading from huggingface. From repo root do

python scripts/download_models.py scripts/downloads --minimal
python scripts/file_server.py

To use local file server with docker pass -e AI_DIFFUSION_DOWNLOAD_URL=http://host.docker.internal:51222 to "docker run" To use local file server in Krita set HOSTMAP=1 environment variable

[Fmstrat]

That's helpful, thank you. For the docker.py script, it throws an error when importing ai_diffusion, stating that the release ZIP wasn't downloaded. I've installed the requirements in the root of the repo, so there is probably some simple step to get it to recognize the folder import that I'm missing (away from PC now).

For SVDQ, I have Nunchaku working in my custom nvidia Dockerfile. I could look to do a PR for that as well, if you would like.

[Acly]

For the docker.py script, it throws an error when importing ai_diffusion, stating that the release ZIP wasn't downloaded.

Likely you're missing a git submodule update

[Fmstrat]

@Acly Yup, that was it. I've got everything building now, but do have an infrastructure question. It seems the Dockerfile has a lot of additional tooling not required to run Krita AI (rclone/Jupyter/etc). I'd like to split this up into the below, and have these auto-build in a GitHub Workflow whenever a formal Release is made so the images are available to pull:

• ghcr.io/Acly/krita-ai-diffusion:comfyui - Just the base required with nvidia support. This would also be tagged latest
• ghcr.io/Acly/krita-ai-diffusion:comfyui-rocm - Just the base required with rocm support
• ghcr.io/Acly/krita-ai-diffusion:comfyui-extended - Additional commonly used tooling + nginx
• ghcr.io/Acly/krita-ai-diffusion:comfyui-rocm-extended - Additional commonly used tooling + nginx

Each of the above would also get a version tag as they were updated on releases. Any objections to this?

[Acly]

You probably saw there's a base version already, I use it to eg. build the cloud images. But what's the use case for making them public? just customization? I think it's more likely people will just find one of the many ComfyUI images and add the few things Krita needs. The existing image is mainly used by people who don't know much about docker and need something that just works.

I believe nginx is required to expose ComfyUI on typical hosters for the websocket-over-http at least. The other stuff also doesn't hurt much all things considered.

Building on GH actions would be nice, so far the images are pushed to https://hub.docker.com/r/aclysia/sd-comfyui-krita/tags though.

[Acly]

I updated models.json to include non-nunchaku version of Kontext, (and fp4 versions for nvidia blackwell)

Also download_models.py should automatically fetch the matching set of models for detected hardware now, see https://github.com/Acly/krita-ai-diffusion/commit/21e7183fbbd21735b3511c616b55f4f2da3f1c0b

Haven't tested this with the docker images yet though.