diffuseR-TS

A collection of Python scripts for exporting TorchScript files compatible with the diffuseR package.

Requirements

Hardware

• NVIDIA GPU (6GB+ VRAM recommended)

Software

• Docker
• NVIDIA Container Toolkit
• NVIDIA Drivers

System Details

Tested on Ubuntu 24.04.2 LTS with NVIDIA GPU (6GB VRAM), but should work on any Unix-like OS with compatible NVIDIA hardware.

Driver Info:

NVIDIA-SMI 570.86.10    Driver Version: 570.86.10    CUDA Version: 12.8

Usage

1. Open docker.sh
2. Build the Docker image
3. Start the Docker container
4. Navigate to the /app directory inside the container (which is bound to the current directory on the host)
5. Run the Python scripts to:
   • Download your chosen diffusion model
   • Load the models
   • Export as TorchScript files (.pt)

The models are downloaded first to avoid re-downloading large (10GB+) files.

Notes

These scripts were developed interactively and may not represent the most elegant solution. The process required some manual intervention and tinkering. Contributions and improvements are welcome via issues or pull requests.