vision_unlearning.unlearner.uce_sd_erase
========================================

.. py:module:: vision_unlearning.unlearner.uce_sd_erase


Classes
-------

.. autoapisummary::

   vision_unlearning.unlearner.uce_sd_erase.ConceptType
   vision_unlearning.unlearner.uce_sd_erase.UCE


Module Contents
---------------

.. py:class:: ConceptType

   Bases: :py:obj:`str`, :py:obj:`enum.Enum`


   Enum representing the type of concept to unlearn.


   .. py:attribute:: Object
      :value: 'object'



   .. py:attribute:: Art
      :value: 'art'



.. py:class:: UCE(**data: Any)

   Bases: :py:obj:`vision_unlearning.unlearner.base.Unlearner`


   Unified Concept Editing for unlearning in Stable Diffusion models.
   Adapted from:
       GitHub: https://github.com/rohitgandikota/unified-concept-editing
       Arxiv: https://arxiv.org/pdf/2308.14761.pdf
       Gandikota, R., Orgad, H., Belinkov, Y., Materzyńska, J., & Bau, D. (2024).
       Unified concept editing in diffusion models. In Proceedings of the IEEE/CVF
       Winter Conference on Applications of Computer Vision (pp. 5111-5120).
   This unlearner do not use LoRA, and do not perform any fine-tuning (instead, it performs a closed-form weight update).


   .. py:attribute:: pretrained_model_name_or_path
      :type:  str
      :value: None



   .. py:attribute:: erase_scale
      :type:  float
      :value: None



   .. py:attribute:: preserve_scale
      :type:  float
      :value: None



   .. py:attribute:: lamb
      :type:  float
      :value: None



   .. py:attribute:: save_entire_model
      :type:  bool
      :value: None



   .. py:attribute:: edit_concepts
      :type:  Optional[str]
      :value: None



   .. py:attribute:: guide_concepts
      :type:  Optional[str]
      :value: None



   .. py:attribute:: preserve_concepts
      :type:  Optional[str]
      :value: None



   .. py:attribute:: concept_type
      :type:  ConceptType
      :value: None



   .. py:attribute:: expand_prompts
      :type:  bool
      :value: True



   .. py:attribute:: final_eval_prompts_forget
      :type:  str | List[str]
      :value: None



   .. py:attribute:: final_eval_prompts_retain
      :type:  str | List[str]
      :value: None



   .. py:attribute:: output_dir
      :type:  str
      :value: None



   .. py:attribute:: device
      :type:  str
      :value: 'cuda:0'



   .. py:attribute:: compute_runtimes
      :type:  bool
      :value: None



   .. py:attribute:: hub_model_id
      :type:  Optional[str]
      :value: None



   .. py:method:: _collect_text_embeddings(pipe: Any, concepts: list[str], device: str, torch_dtype: torch.dtype) -> dict[str, torch.Tensor]

      Return dict {concept: last_token_embedding}.



   .. py:method:: _collect_guide_outputs(concepts: list[str], embeds: dict[str, torch.Tensor], modules: list[torch.nn.Module]) -> dict[str, list[torch.Tensor]]

      Collect cross-attention outputs for guide/preserve concepts.



   .. py:method:: _update_weights(original_modules: list[torch.nn.Module], erase_embeds: dict[str, torch.Tensor], guide_outputs: dict[str, list[torch.Tensor]], edit_concepts: list[str], guide_concepts: list[str], preserve_concepts: list[str], erase_scale: float, preserve_scale: float, lamb: float, device: str, torch_dtype: torch.dtype) -> list[torch.nn.Module]

      Apply the UCE weight update to each module and return new modules.



   .. py:method:: _save_uce_weights(uce_modules: list[torch.nn.Module], uce_module_names: list[str]) -> None

      Save updated module weights to a safetensors file.



   .. py:method:: train() -> List[huggingface_hub.repocard_data.EvalResult]

      Main UCE training and concept erasure logic.



   .. py:method:: get_pipeline_from_modified_weights(pretrained_model_name_or_path: str, device: str | torch.device, output_dir: str) -> diffusers.DiffusionPipeline
      :staticmethod:



   .. py:method:: evaluate() -> Tuple[List[huggingface_hub.repocard_data.EvalResult], Dict[str, PIL.Image.Image]]