integrate diffusion model FE-BE-GPU

backend/src/constants/algorithmsConstants.js

@@ -119,21 +119,21 @@ const TRAIN_ALGORITHMS = {
 };
 
 const MULTIFLOW_ALGORITHMS = {
-	PTQ4DIT_QUANT_SAMPLE: {
+	MULTIFLOW_PRUNE: {
 		path: MULTIFLOW_PATH,
 		type: ALGORITHM_TYPES.MULTIFLOW,
 		fileName: 'prune.py'
 	}
 };
 
 const DIFFUSION_MODEL_ALGORITHMS = {
-	PTQ4DIT_GET_CALIBRATION_SET: {
-		path: DIFFUSION_MODEL_PATH / 'ptq4dit/',
+	PTQ4DiT_GET_CALIBRATION_SET: {
+		path: `${DIFFUSION_MODEL_PATH}PTQ4DiT/`,
 		type: ALGORITHM_TYPES.DIFFUSION_MODEL,
 		fileName: 'get_calibration_set.py'
 	},
-	PTQ4DIT_QUANT_SAMPLE: {
-		path: DIFFUSION_MODEL_PATH / 'ptq4dit/',
+	PTQ4DiT_QUANT_SAMPLE: {
+		path: `${DIFFUSION_MODEL_PATH}PTQ4DiT/`,
 		type: ALGORITHM_TYPES.DIFFUSION_MODEL,
 		fileName: 'quant_sample.py'
 	}

backend/src/constants/parametersConstants.js

@@ -1528,7 +1528,7 @@ const MULTIFLOW_PARAMETERS = {
 };
 
 const DIFFUSION_MODEL_PARAMETERS = {
-	PTQ4DIT_GET_CALIBRATION_SET: [
+	PTQ4DiT_GET_CALIBRATION_SET: [
 		{
 			argName: 'model',
 			defaultValue: 'DiT-XL/2',
@@ -1608,7 +1608,7 @@ const DIFFUSION_MODEL_PARAMETERS = {
 		},
 		{
 			argName: 'ckpt',
-			defaultValue: null,
+			defaultValue: '',
 			inputType: 'text',
 			label: 'Checkpoint Path',
 			placeholder: 'Enter checkpoint path',
@@ -1631,7 +1631,7 @@ const DIFFUSION_MODEL_PARAMETERS = {
 			help: 'Name of the file where generated logs will be saved.'
 		}
 	],
-	PTQ4DIT_QUANT_SAMPLE: [
+	PTQ4DiT_QUANT_SAMPLE: [
 		{
 			argName: 'model',
 			defaultValue: 'DiT-XL/2',
@@ -1711,7 +1711,7 @@ const DIFFUSION_MODEL_PARAMETERS = {
 		},
 		{
 			argName: 'ckpt',
-			defaultValue: null,
+			defaultValue: '',
 			inputType: 'text',
 			label: 'Checkpoint Path',
 			placeholder: 'Enter checkpoint path',

backend/src/router/scriptsRouter.js

@@ -69,6 +69,8 @@ const buildCommand = (scriptDetails, args) => {
 
 	if (scriptDetails.type === ALGORITHM_TYPES.MULTIFLOW) {
 		return `${condaActivate} && cd ${basePath}/multiflow && python3 ${scriptDetails.fileName} ${args}`;
+	} else if (scriptDetails.type === ALGORITHM_TYPES.DIFFUSION_MODEL) {
+		return `${condaActivate} && cd ${basePath}/diffusion_model/PTQ4DiT && python3 ${scriptDetails.fileName} ${args}`;
 	}
 
 	return `${condaActivate} && python3 ${scriptPath} ${args}`;
@@ -99,6 +101,7 @@ const getCurrentOrLastActiveScriptDetails = (req, res) => {
 
 const runScript = async (req, res) => {
 	const { alg = '', params = {} } = req.body;
+
 	const scriptDetails = ALGORITHMS[alg];
 
 	if (!scriptDetails) {

docs/configure-diffusion-model.md

@@ -0,0 +1,164 @@
+# PTQ4DiT Configuration Guide
+
+This guide provides detailed steps to configure and set up the **PTQ4DiT (Post-training Quantization for Diffusion Transformers)** project.
+
+## Step-by-Step Configuration
+
+### Step 1: Navigate to the PTQ4DiT Directory
+
+Change your current directory to the `PTQ4DiT` project folder.
+
+```bash
+cd ~/machine_learning_core/diffusion_model/PTQ4DiT
+```
+
+### Step 2: Activate the Conda Environment
+
+Activate your existing Conda environment named `modelsmith`.
+
+```bash
+conda activate modelsmith
+```
+
+### Step 3: Update the Conda Environment
+
+Ensure that all necessary dependencies for PTQ4DiT are installed by updating the Conda environment using the provided `environment.yml` file.
+
+```bash
+conda env update --file environment.yml --name modelsmith
+```
+
+### Step 4: Create Necessary Directories
+
+Create directories for calibration data and pre-trained model checkpoints.
+
+```bash
+mkdir calib
+mkdir -p pretrained_models
+```
+
+### Step 5: Download Pre-trained DiT Checkpoints
+
+Navigate to the `pretrained_models` directory and download the required DiT checkpoints.
+
+```bash
+cd pretrained_models
+
+# Download 256×256 DiT checkpoint
+wget https://dl.fbaipublicfiles.com/DiT/models/DiT-XL-2-256x256.pt
+
+# Download 512×512 DiT checkpoint
+wget https://dl.fbaipublicfiles.com/DiT/models/DiT-XL-2-512x512.pt
+```
+
+**Note:** If you encounter any issues with `wget`, you can use `curl` as an alternative.
+
+```bash
+# Using curl to download 256×256 DiT checkpoint
+curl -L https://dl.fbaipublicfiles.com/DiT/models/DiT-XL-2-256x256.pt -o DiT-XL-2-256x256.pt
+
+# Using curl to download 512×512 DiT checkpoint
+curl -L https://dl.fbaipublicfiles.com/DiT/models/DiT-XL-2-512x512.pt -o DiT-XL-2-512x512.pt
+```
+
+### Step 6: Generate Calibration Data (can be done from the Diffusion Model page)
+
+Return to the `PTQ4DiT` directory and generate the calibration datasets. Replace `256` with `512` in the `--image-size` parameter if you wish to generate calibration data for 512×512 images.
+
+```bash
+cd ~/machine_learning_core/diffusion_model/PTQ4DiT
+
+python get_calibration_set.py --model DiT-XL/2 --image-size 256 \
+--ckpt pretrained_models/DiT-XL-2-256x256.pt \
+--num-sampling-steps 50 \
+--outdir calib/ --filename imagenet_DiT-256_sample4000_50steps_allst.pt \
+--cfg-scale 1.5 --seed 1
+```
+
+**Parameters Explanation:**
+
+- `--model`: Specifies the model variant (e.g., `DiT-XL/2`).
+- `--image-size`: Image resolution (256 or 512).
+- `--ckpt`: Path to the pre-trained checkpoint.
+- `--num-sampling-steps`: Number of sampling steps.
+- `--outdir`: Output directory for calibration data.
+- `--filename`: Name of the output calibration file.
+- `--cfg-scale`: Configuration scale parameter.
+- `--seed`: Random seed for reproducibility.
+
+### Step 7: Perform Quantization (can be done from the Diffusion Model page)
+
+After generating the calibration data, perform quantization using the `quant_sample.py` script. You can specify different bit-widths and sampling steps as needed.
+
+#### Example 1: Quantize to W8A8 for 256×256 Images
+
+```bash
+python quant_sample.py --model DiT-XL/2 --image-size 256 \
+--ckpt pretrained_models/DiT-XL-2-256x256.pt \
+--num-sampling-steps 50 \
+--weight_bit 8 --act_bit 8 --cali_st 25 --cali_n 64 --cali_batch_size 32 --sm_abit 8 \
+--cali_data_path calib/imagenet_DiT-256_sample4000_50steps_allst.pt --outdir output/ \
+--cfg-scale 1.5 --seed 1 --ptq --recon
+```
+
+#### Example 2: Quantize to W4A8 for 512×512 Images
+
+```bash
+python quant_sample.py --model DiT-XL/2 --image-size 512 \
+--ckpt pretrained_models/DiT-XL-2-512x512.pt \
+--num-sampling-steps 50 \
+--weight_bit 4 --act_bit 8 --cali_st 10 --cali_n 128 --cali_batch_size 16 --sm_abit 8 \
+--cali_data_path calib/imagenet_DiT-512_sample4000_50steps_allst.pt --outdir output/ \
+--cfg-scale 1.5 --seed 1 --ptq --recon
+```
+
+**Parameters Explanation:**
+
+- `--weight_bit` & `--act_bit`: Specify the bit-width for weights and activations.
+- `--cali_st`, `--cali_n`, `--cali_batch_size`: Calibration parameters.
+- `--sm_abit`: Additional quantization parameter.
+- `--cali_data_path`: Path to the calibration data.
+- `--outdir`: Output directory for quantized models.
+- `--ptq`: Enables post-training quantization.
+- `--recon`: Enables reconstruction during quantization.
+
+### Step 8: Perform Inference (can be done from the Diffusion Model page)
+
+Use the quantized models for inference by resuming the calibration and running the inference algorithm.
+
+#### Example 1: Inference with W8A8 for 256×256 Images
+
+```bash
+python quant_sample.py --model DiT-XL/2 --image-size 256 \
+--ckpt pretrained_models/DiT-XL-2-256x256.pt \
+--num-sampling-steps 50 \
+--weight_bit 8 --act_bit 8 --cali_st 25 --cali_n 64 --cali_batch_size 32 --sm_abit 8 \
+--cali_data_path calib/imagenet_DiT-256_sample4000_50steps_allst.pt --outdir output/ \
+--cfg-scale 1.5 --seed 1 \
+--resume --cali_ckpt output/256_88_50/ckpt.pth \
+--ptq \
+--inference --n_c 10
+```
+
+#### Example 2: Inference with W4A8 for 512×512 Images
+
+```bash
+python quant_sample.py --model DiT-XL/2 --image-size 512 \
+--ckpt pretrained_models/DiT-XL-2-512x512.pt \
+--num-sampling-steps 50 \
+--weight_bit 4 --act_bit 8 --cali_st 10 --cali_n 128 --cali_batch_size 16 --sm_abit 8 \
+--cali_data_path calib/imagenet_DiT-512_sample4000_50steps_allst.pt --outdir output/ \
+--cfg-scale 1.5 --seed 1 \
+--resume --cali_ckpt output/512_48_50/ckpt.pth \
+--ptq \
+--inference --n_c 5
+```
+
+## Troubleshooting
+
+If you encounter issues, ensure:
+
+- **Correct paths:** File paths exist and are accessible.
+- **Matching versions:** PyTorch and diffusers are compatible.
+
+---

docs/configure-modelsmith-on-local-machine-docker.md

@@ -85,6 +85,12 @@ Still inside the backend container, configure Multiflow by following the instruc
 
 These configurations are necessary to enable 100% of the application's functionalities. If skipped, only machine learning, quantization, and pruning features will work.
 
+### Step 7: Configure Diffusion Model (PTQ4DiT) Python Project
+
+Still inside the backend container, configure Diffusion Model by following the instructions in the [Diffusion Model Configuration Guide](configure-diffusion-model.md):
+
+These configurations are necessary to enable 100% of the application's functionalities. If skipped, only machine learning, quantization, and pruning features will work.
+
 ### Backend Container Initialization
 
 The backend container will perform the following initialization steps:

docs/configure-modelsmith-on-vm-docker.md

@@ -118,6 +118,12 @@ To enable multiflow functionalities, please refer to the [Multiflow Configuratio
 
 These configurations are necessary to enable 100% of the application's functionalities. If skipped, only machine learning, quantization, and pruning features will work.
 
+### Step 6: Configure Diffusion Model (PTQ4DiT) Python Project
+
+Still inside the backend container, configure Diffusion Model by following the instructions in the [Diffusion Model Configuration Guide](configure-diffusion-model.md):
+
+These configurations are necessary to enable 100% of the application's functionalities. If skipped, only machine learning, quantization, and pruning features will work.
+
 ### Backend Container Initialization
 
 The backend container will perform the following initialization steps:

docs/configure-modelsmith-on-vm.md

@@ -98,6 +98,10 @@ conda --version
 
    Please refer to the [Multiflow Configuration Guide](configure-multiflow.md) for detailed instructions.
 
+8. **Configure Diffusion Model Python Project**:
+
+   Please refer to the [Diffusion Model Configuration Guide](configure-diffusion-model.md) for detailed instructions.
+
 ## Phase 2: Configure the Environment
 
 After configuring the Python environment, the next step involves setting up the project environment. This is done by executing the `setup_environment.sh` script located in the `backend` directory. The script performs several tasks, including setting up the `.env` file for the backend server.

frontend/src/app/modules/model-compression/models/enums/algorithms.enum.ts

@@ -51,8 +51,8 @@ export enum MultiflowAlgorithmsEnum {
 }
 
 export enum DiffusionModelAlgorithmsEnum {
-	PTQ4DIT_GET_CALIBRATION_SET = 'PTQ4DIT_GET_CALIBRATION_SET',
-	PTQ4DIT_QUANT_SAMPLE = 'PTQ4DIT_QUANT_SAMPLE'
+	PTQ4DIT_GET_CALIBRATION_SET = 'PTQ4DiT_GET_CALIBRATION_SET',
+	PTQ4DIT_QUANT_SAMPLE = 'PTQ4DiT_QUANT_SAMPLE'
 }
 
 export enum PruningAlgorithmsEnum {