LDG-2: Diffusion Model

[henrikfo]

Description

Summary of changes

• This pull request introduces the DiffMI attack with an example for CelebA. The main changes include adding a new input handler for the DiffMI attack, the Diff-Mi attack itself and extending the audit configuration with new parameters for the DiffMI attack.

DiffMI Attack Implementation:

• Instead of implementing a new CelebA_InputHandler in a file called celebA_diffmi_handler.py for the diffusion model training, a specific file called train_utils.py is created in attacks/utils/diffmi_utils/ with training specific to the DiffMI attack. The reason being that the training procedure and all its functions is very complex.

Configuration Updates for DiffMI:

• Extended audit.yaml to include a new diffmi attack section with parameters for fine-tuning, preprocessing, pretraining, and attack-specific settings.

Evaluation Pipeline:

• Added a MinvResults class.

Minimal test coverage

• Diff-MI diffusion handler tests for built-in/default diffusion loading.
• tests for custom module_path loading behavior.
• checkpoint completion-marker tests to avoid using interrupted pretrain/fine-tune checkpoints.

How Has This Been Tested?

The attack has been tested with and without minibatch for H100 and 2080ti respectively. Fine-tuning on a 2080ti is not recommended since only a batch_size of 1 is possible.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[TheColdIce]

I have gone through all of the code now @henrikfo. Nice work!

• I have made some comments that needs to be fix/clarified.
• I think I have flagged all code that is not used.
• I have not flagged ruff errors. I also noticed some inconsistency in the docstring format, I dont know if ruff will flag this.
• I think also the wiki needs to be updated with how to run the attack. There is some params, use_fp16 for example, that is not clarified in the audit.yaml. Alternatively, the audit file can be flashed out.

If the comments are resolved, ruff checks are fixed and there is some clarification regarding config params, I think we can merge with main.

[henrikfo]

• I have made some comments that needs to be fix/clarified.

I believe that all of them are now resolved or responded to!

• I think I have flagged all code that is not used.

Those files/line have been removed!

• I have not flagged ruff errors. I also noticed some inconsistency in the docstring format, I dont know if ruff will flag this.

All ruff checked have passed and the inconsistencies in docstrings should be resolved aswell!

• I think also the wiki needs to be updated with how to run the attack. There is some params, use_fp16 for example, that is not clarified in the audit.yaml. Alternatively, the audit file can be flashed out.

Yes, the use_fp16 bug is fixed and I will create an new Issue about creating a comprehensive wiki for the attack!

If the comments are resolved, ruff checks are fixed and there is some clarification regarding config params, I think we can merge with main.

Great!

[TheColdIce]

• Commented on a couple of minor fixes.

• With license update there is also headers needed at the top of every file. I suggest merge with main, then add the header to files that are missing:

#
# Copyright 2023-2026 Lindholmen Science Park AB
# SPDX-License-Identifier: Apache-2.0
#

• The NOTICE file should also be updated with refs to files that use externel sources, like:

  • leakpro/attacks/utils/diff_mi/losses.py
  • leakpro/attacks/utils/diff_mi/unet.py

• And the licenses of packages like kornia, blobfile, lpips, robustness and mpi4py need to be checked that they are compatible with apache-2.0

[henrikfo]

• Commented on a couple of minor fixes.

Fixed!

• With license update there is also headers needed at the top of every file. I suggest merge with main, then add the header to files that are missing:

#
# Copyright 2023-2026 Lindholmen Science Park AB
# SPDX-License-Identifier: Apache-2.0
#

Done!

• The NOTICE file should also be updated with refs to files that use externel sources, like:

  • leakpro/attacks/utils/diff_mi/losses.py
  • leakpro/attacks/utils/diff_mi/unet.py

• And the licenses of packages like kornia, blobfile, lpips, robustness and mpi4py need to be checked that they are compatible with apache-2.0

Done!

[TheColdIce]

celebA_plgmi_handler.py

the evaluation() function needs to be named eval() otherwise leakpro throws an error.

[TheColdIce]

I got this error when running the examples/minv/celebA/main_celebA_minv.ipynb notebook:

AttributeError                            Traceback (most recent call last)
Cell In[7], [line 10](vscode-notebook-cell:?execution_count=7&line=10)
      6 # Initialize the LeakPro object
      7 leakpro = LeakPro(CelebA_InputHandler, config_path)
      8 
      9 # Run the audit
---> [10](vscode-notebook-cell:?execution_count=7&line=10) results = leakpro.run_audit()

File ~/LeakPro/leakpro/leakpro.py:174, in LeakPro.run_audit(self, create_pdf, use_optuna)
    171 def run_audit(self:Self, create_pdf: bool = False, use_optuna: bool = False) -> list[Any]:
    172     """Run the audit."""
--> [174](https://vscode-remote+ssh-002dremote-002b172-002e25-002e17-002e197.vscode-resource.vscode-cdn.net/home/edvin/LeakPro/examples/minv/celebA/~/LeakPro/leakpro/leakpro.py:174)     audit_results = self.attack_scheduler.run_attacks(use_optuna=use_optuna)
    175     results = [entry["result_object"] for entry in audit_results]
    177     if create_pdf:

File ~/LeakPro/leakpro/attacks/attack_scheduler.py:126, in AttackScheduler.run_attacks(self, use_optuna)
    124     result = attack_obj.run_attack()
    125     logger.info(f"Saving results for attack: {attack_type} to {self.report_dir}")
--> [126](https://vscode-remote+ssh-002dremote-002b172-002e25-002e17-002e197.vscode-resource.vscode-cdn.net/home/edvin/LeakPro/examples/minv/celebA/~/LeakPro/leakpro/attacks/attack_scheduler.py:126)     result.save(attack_obj = attack_obj, output_dir = self.output_dir)
    127     results.append({"attack_type": attack_type, "attack_object": attack_obj, "result_object": result})
    129 return results

AttributeError: 'ImageMetrics' object has no attribute 'save'

I used ResNet18 and the plgmi attack.

[TheColdIce]

Ping @henrikfo
Did you want to fix this comment in this PR or du you want me to create a issue?

[henrikfo]

Ping @henrikfo Did you want to fix this comment in this PR or du you want me to create a issue?

No, I've fixed it already! I'll push updates in just a minute or two!