DiffusionLens Tutorial

[AdamBelfki3]

[ericluo04]

Thank you for opening this issue! On the topic, I was kindly wondering if the below points could be covered:

• How to do interventions on not only all timesteps (with model.all() as in the colab from #362) but also on specific timesteps.
• How to extract activations from specific timesteps (e.g., jointly save the residual stream from the 24th SD3.5 layer, for the 1st, 3rd, and 6th timesteps).
• How to do both of the above with img2img, as in #417.

Any tips on these points in the meantime would be much appreciated.

Look forward to using nnsight in my own research and am eager to see the tutorial whenever it's ready - thank you!

[AdamBelfki3]

Hi @ericluo04 - happy to see your excitement for using nnsight :)

For both your questions (1) and (2), you can replace the with module.all() syntax with with module.iter[start_idx:end_idx]: to access one or more specific generation time steps.

For img2img generation (3) we will integrate the PR from @Kmoneal and make proper documentation for it.

Let me know if you have more questions in the meantime.

[ericluo04]

Thanks again for your response @AdamBelfki3! This tutorial has already been super helpful. :)

Two more follow-up questions:

1. Is it possible to add an already existing LoRa to an nnsight model? For example, for an img2img pipeline, I can do the below. Is there something analogous with nnsight?

pipe = AutoPipelineForImage2Image.from_pretrained("stabilityai/stable-diffusion-3.5-large", torch_dtype=torch.bfloat16)
pipe = pipe.to("cuda")

pipe.load_lora_weights("/user/checkpoint-25400", 
                       weight_name="pytorch_lora_weights.safetensors")
pipe.fuse_lora(lora_scale=.75)
pipe.unload_lora_weights()

2. This may be out-of-scope, but I've had some difficulty implementing nnsight with img2img pipelines (e.g., the PR from Kmoneal, who I have also pinged in #417). For example, I sometimes come across the following error: NNsightError: Img2ImgDiffusionModel._generate() missing 1 required positional argument: 'prepared_inputs' I understand it is not supported right now, but any initial advice or pointers on how to even just setup img2img with nnsight would be much appreciated!

[ericluo04]

Figured it out!

For (1), I added the following to class Img2ImgDiffusionModel(RemoteableMixin): from #417. Disclaimer, I had Gemini 2.5 Pro help write this code - I did check that this implementation works for stabilityai/stable-diffusion-3.5-large, with adjusting lora_scale values appropriately changing model behavior. In particular, I checked that setting lora_scale = 0 produces the same outputs as the base model (without LoRa) for a given seed.

    def load_lora_weights(self, pretrained_model_name_or_path, 
                         weight_name: str = None, 
                         adapter_name: str = "default", 
                         **kwargs):
        """
        Load LoRA weights into the model.
        
        Args:
            pretrained_model_name_or_path: Path to the LoRA weights or model directory
            weight_name: Name of the weight file (e.g., "pytorch_lora_weights.safetensors")
            adapter_name: Name to identify this specific LoRA adapter
            **kwargs: Additional arguments to pass to the underlying load_lora_weights method
        """
        return self._model.pipeline.load_lora_weights(
            pretrained_model_name_or_path,
            weight_name=weight_name,
            adapter_name=adapter_name,
            **kwargs
        )
    
    def fuse_lora(self, lora_scale: float = 1.0, adapter_names: Optional[List[str]] = None):
        """
        Fuse the LoRA weights into the base model weights for inference.
        
        Args:
            lora_scale: Scaling factor for the LoRA weights
            adapter_names: List of adapter names to fuse. If None, fuses all loaded adapters.
        """
        if hasattr(self._model.pipeline, "fuse_lora"):
            return self._model.pipeline.fuse_lora(lora_scale=lora_scale)
        else:
            # For older versions of diffusers
            return self._model.pipeline.set_adapters("default", adapter_weights=lora_scale)
    
    def unfuse_lora(self):
        """
        Unfuse the LoRA weights from the base model weights.
        """
        return self._model.pipeline.unfuse_lora()
    
    def unload_lora_weights(self):
        """
        Unload LoRA weights from the model.
        
        Note: The diffusers API for unload_lora_weights does not accept adapter_names parameter.
        """
        return self._model.pipeline.unload_lora_weights()
    
    def set_lora_scale(self, lora_scale: float, adapter_names: Optional[List[str]] = None):
        """
        Set the scaling factor for the LoRA weights.
        
        Args:
            lora_scale: Scaling factor for the LoRA weights
            adapter_names: List of adapter names to scale. If None, scales all loaded adapters.
        """
        if hasattr(self._model.pipeline, "set_adapters_scale"):
            return self._model.pipeline.set_adapters_scale(lora_scale, adapter_names=adapter_names)
        else:
            raise AttributeError("This pipeline does not support setting LoRA scale dynamically.")

Loading a pre-existing LoRa then simply involves the following:

# initialize the img2img diffusion model with bfloat16 precision and dispatch=True for nnsight
pipe = Img2ImgDiffusionModel("stabilityai/stable-diffusion-3.5-large", torch_dtype=torch.bfloat16, dispatch=True)

# move model to GPU
pipe = pipe.to('cuda')

# load and fuse LoRA weights
pipe.load_lora_weights("/user/checkpoint-25400", 
                      weight_name="pytorch_lora_weights.safetensors")
pipe.fuse_lora(lora_scale=.75)
pipe.unload_lora_weights()

For (2), I included a sample implementation in #417.

Thanks again for all your work on nnsight - it is really an amazing package!! Can't believe how much cleaner my code has gotten. :)

[AdamBelfki3]

@ericluo04 https://nnsight.net/notebooks/tutorials/diffusion_lens/