Skip to content

Base Classes for Diffusers Integration

The base callback classes for integration with Diffusers.

BaseDiffusersCallback

Bases: ABC

Base callback for 🧨 Diffusers logging the results of a DiffusionPipeline generation to Weights & Biases.

Parameters:

Name Type Description Default
pipeline DiffusionPipeline

The DiffusionPipeline from diffusers.

 required
prompt Union[str, List[str]]

The prompt or prompts to guide the image generation.

 required
wandb_project Optional[str]

The name of the project where you're sending the new run. The project is not necessary to be specified unless the run has automatically been initiatlized before the callback is defined.

 required
wandb_entity Optional[str]

An entity is a username or team name where you're sending runs. This entity must exist before you can send runs there, so make sure to create your account or team in the UI before starting to log runs. If you don't specify an entity, the run will be sent to your default entity, which is usually your username. Change your default entity in your settings under "default location to create new projects".

 None
weave_mode bool

Whether to use log to a weave board instead of W&B dashboard or not. The weave mode logs the configs, generated images and timestamp in a StreamTable instead of a wandb.Table and does not require a wandb run to be initialized in order to start logging. This makes it possible to log muliple generations without having to initialize or terminate runs. Note that the parameter wandb_entity must be explicitly specified in order to use weave mode.

 False
num_inference_steps int

The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference.

 50
num_images_per_prompt Optional[int]

The number of images to generate per prompt.

 1
negative_prompt Optional[Union[str, List[str]]]

The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if guidance_scale is less than 1).

 None
configs Optional[Dict]

Additional configs for the experiment you want to sync, for example, for example, seed could be a good config to be passed here.

 None
Source code in wandb_addons/diffusers/callbacks/base/base_callback.py 
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
class BaseDiffusersCallback(ABC):
    """Base callback for [🧨 Diffusers](https://huggingface.co/docs/diffusers/index)
    logging the results of a
    [`DiffusionPipeline`](https://github.com/huggingface/diffusers/blob/v0.21.0/src/diffusers/pipelines/pipeline_utils.py#L480)
    generation to Weights & Biases.

    Arguments:
        pipeline (diffusers.DiffusionPipeline): The `DiffusionPipeline` from
            `diffusers`.
        prompt (Union[str, List[str]]): The prompt or prompts to guide the image
            generation.
        wandb_project (Optional[str]): The name of the project where you're sending
            the new run. The project is not necessary to be specified unless the run
            has automatically been initiatlized before the callback is defined.
        wandb_entity (Optional[str]): An entity is a username or team name where
            you're sending runs. This entity must exist before you can send runs there,
            so make sure to create your account or team in the UI before starting to
            log runs. If you don't specify an entity, the run will be sent to your
            default entity, which is usually your username. Change your default entity
            in [your settings](https://wandb.ai/settings) under "default location to
            create new projects".
        weave_mode (bool): Whether to use log to a
            [weave board](https://docs.wandb.ai/guides/weave) instead of W&B dashboard
            or not. The weave mode logs the configs, generated images and timestamp in a
            [`StreamTable`](https://docs.wandb.ai/guides/weave/streamtable) instead of a
            `wandb.Table` and does not require a wandb run to be initialized in order to
            start logging. This makes it possible to log muliple generations without
            having to initialize or terminate runs. Note that the parameter
            `wandb_entity` must be explicitly specified in order to use weave mode.
        num_inference_steps (int): The number of denoising steps. More denoising steps
            usually lead to a higher quality image at the expense of slower inference.
        num_images_per_prompt (Optional[int]): The number of images to generate per
            prompt.
        negative_prompt (Optional[Union[str, List[str]]]): The prompt or prompts not
            to guide the image generation. Ignored when not using guidance
            (i.e., ignored if `guidance_scale` is less than `1`).
        configs (Optional[Dict]): Additional configs for the experiment you want to
            sync, for example, for example, `seed` could be a good config to be passed
            here.
    """

    def __init__(
        self,
        pipeline: DiffusionPipeline,
        prompt: Union[str, List[str]],
        wandb_project: str,
        wandb_entity: Optional[str] = None,
        weave_mode: bool = False,
        num_inference_steps: int = 50,
        num_images_per_prompt: Optional[int] = 1,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        configs: Optional[Dict] = None,
        **kwargs,
    ) -> None:
        super().__init__()
        self.pipeline = pipeline
        self.prompt = prompt
        self.wandb_project = wandb_project
        self.wandb_entity = wandb_entity
        self.weave_mode = weave_mode
        self.num_inference_steps = num_inference_steps
        self.num_images_per_prompt = num_images_per_prompt
        self.negative_prompt = negative_prompt
        self.configs = configs
        self.wandb_table = None
        self.table_row = []
        self.starting_step = 1
        self.log_step = num_inference_steps
        self.job_type = "text-to-image"
        self.table_name = "Text-To-Image"
        self.initialize_wandb(wandb_project, wandb_entity)
        self.build_wandb_table()

    def update_configs(self) -> None:
        """Update the configs as a state of the callback. This function is called inside
        `initialize_wandb()`.
        """
        pipeline_configs = dict(self.pipeline.config)
        pipeline_configs["scheduler"] = list(pipeline_configs["scheduler"])
        pipeline_configs["scheduler"][1] = dict(self.pipeline.scheduler.config)
        additional_configs = {
            "prompt": self.prompt,
            "negative_prompt": self.negative_prompt
            if self.negative_prompt is not None
            else "",
            "num_inference_steps": self.num_inference_steps,
            "num_images_per_prompt": self.num_images_per_prompt,
            "pipeline": pipeline_configs,
        }
        self.configs = (
            {**self.configs, **additional_configs}
            if self.configs is not None
            else additional_configs
        )

    def initialize_wandb(self, wandb_project: str, wandb_entity: str) -> None:
        """Initializes the wandb run if not already initialized. If `weave_mode` is
        `True`, then a [StreamTable](https://docs.wandb.ai/guides/weave/streamtable) is
        initialized instead of a wandb run. This function is called automatically when
        the callback is initialized.

        Arguments:
            wandb_project (str): The name of the W&B project.
            wandb_entity (str): The name of the W&B entity.
        """
        self.update_configs()
        if self.weave_mode:
            if self.wandb_entity is None:
                wandb.termerror(
                    "The parameter wandb_entity must be provided when weave_mode is enabled."
                )
            else:
                self.stream_table = StreamTable(
                    f"{self.wandb_entity}/{self.wandb_project}/{self.table_name}"
                )
                self.table_row = []
        else:
            if wandb.run is None:
                if wandb_project is not None:
                    wandb.init(
                        project=wandb_project,
                        entity=wandb_entity,
                        job_type=self.job_type,
                        config=self.configs,
                    )
                else:
                    wandb.termerror("The parameter wandb_project must be provided.")

    def build_wandb_table(self) -> None:
        """Specifies the columns of the wandb table if not in weave mode. This function
        is called automatically when the callback is initialized.
        """
        self.table_columns = [
            "Prompt",
            "Negative-Prompt",
            "Generated-Image",
            "Image-Size",
        ]

    @abstractmethod
    def generate(self, latents: torch.FloatTensor) -> List:
        """Generate images from latents. This method must be implemented in the child
        class."""
        pass

    def populate_table_row(
        self, prompt: str, negative_prompt: str, image: Image
    ) -> None:
        """Populates the table row with the prompt, negative prompt, the generated
        image, and the configs.

        Arguments:
            prompt (str): The prompt to guide the image generation.
            negative_prompt (str): The prompt not to guide the image generation.
            image (Image): The generated image.
        """
        width, height = image.size
        if self.weave_mode:
            self.table_row += [
                {
                    "Generated-Image": image,
                    "Image-Size": {"Width": width, "Height": height},
                    "Configs": self.configs,
                }
            ]
        else:
            self.table_row = [
                prompt,
                negative_prompt if negative_prompt is not None else "",
                wandb.Image(image),
                {"Width": width, "Height": height},
            ]

    def at_initial_step(self):
        """A function that will be called at the initial step of the denoising loop
        during inference."""
        if not self.weave_mode:
            self.wandb_table = wandb.Table(columns=self.table_columns)

    def __call__(
        self,
        step: int,
        timestep: int,
        latents: torch.FloatTensor,
        end_experiment: bool = True,
    ):
        """A function that will be called every `callback_steps` steps during
        inference with the `diffusers.DiffusionPipeline`.

        Arguments:
            step (int): The current step of the inference.
            timestep (int): The current timestep of the inference.
            latents (torch.FloatTensor): The latent tensor used to generate the image.
            end_experiment (bool): Whether to end the experiment automatically or not
                after the pipeline is executed.
        """
        if step == self.starting_step:
            self.at_initial_step()
        if step == self.log_step:
            images = self.generate(latents)
            prompt_logging = (
                self.prompt if isinstance(self.prompt, list) else [self.prompt]
            )
            negative_prompt_logging = (
                self.negative_prompt
                if isinstance(self.negative_prompt, list)
                else [self.negative_prompt] * len(prompt_logging)
            )
            images = chunkify(images, len(prompt_logging))
            for idx in range(len(prompt_logging)):
                for image in images[idx]:
                    self.populate_table_row(
                        prompt_logging[idx], negative_prompt_logging[idx], image
                    )
                    if not self.weave_mode:
                        self.wandb_table.add_data(*self.table_row)
                        wandb.log({"Generated-Images": wandb.Image(image)})
            if end_experiment:
                self.end_experiment()

    def end_experiment(self):
        """Ends the experiment. This function is called automatically at the end of
        `__call__` the parameter `end_experiment` is set to `True`.
        """
        if self.weave_mode:
            self.stream_table.log(self.table_row)
            self.stream_table.finish()
        elif wandb.run is not None:
            wandb.log({self.table_name: self.wandb_table})
            wandb.finish()

__call__(step, timestep, latents, end_experiment=True)

A function that will be called every callback_steps steps during inference with the diffusers.DiffusionPipeline.

Parameters:

Name Type Description Default
step int

The current step of the inference.

 required
timestep int

The current timestep of the inference.

 required
latents FloatTensor

The latent tensor used to generate the image.

 required
end_experiment bool

Whether to end the experiment automatically or not after the pipeline is executed.

 True
Source code in wandb_addons/diffusers/callbacks/base/base_callback.py 
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
def __call__(
    self,
    step: int,
    timestep: int,
    latents: torch.FloatTensor,
    end_experiment: bool = True,
):
    """A function that will be called every `callback_steps` steps during
    inference with the `diffusers.DiffusionPipeline`.

    Arguments:
        step (int): The current step of the inference.
        timestep (int): The current timestep of the inference.
        latents (torch.FloatTensor): The latent tensor used to generate the image.
        end_experiment (bool): Whether to end the experiment automatically or not
            after the pipeline is executed.
    """
    if step == self.starting_step:
        self.at_initial_step()
    if step == self.log_step:
        images = self.generate(latents)
        prompt_logging = (
            self.prompt if isinstance(self.prompt, list) else [self.prompt]
        )
        negative_prompt_logging = (
            self.negative_prompt
            if isinstance(self.negative_prompt, list)
            else [self.negative_prompt] * len(prompt_logging)
        )
        images = chunkify(images, len(prompt_logging))
        for idx in range(len(prompt_logging)):
            for image in images[idx]:
                self.populate_table_row(
                    prompt_logging[idx], negative_prompt_logging[idx], image
                )
                if not self.weave_mode:
                    self.wandb_table.add_data(*self.table_row)
                    wandb.log({"Generated-Images": wandb.Image(image)})
        if end_experiment:
            self.end_experiment()

at_initial_step()

A function that will be called at the initial step of the denoising loop during inference.

Source code in wandb_addons/diffusers/callbacks/base/base_callback.py 
186
187
188
189
190
def at_initial_step(self):
    """A function that will be called at the initial step of the denoising loop
    during inference."""
    if not self.weave_mode:
        self.wandb_table = wandb.Table(columns=self.table_columns)

build_wandb_table()

Specifies the columns of the wandb table if not in weave mode. This function is called automatically when the callback is initialized.

Source code in wandb_addons/diffusers/callbacks/base/base_callback.py 
141
142
143
144
145
146
147
148
149
150
def build_wandb_table(self) -> None:
    """Specifies the columns of the wandb table if not in weave mode. This function
    is called automatically when the callback is initialized.
    """
    self.table_columns = [
        "Prompt",
        "Negative-Prompt",
        "Generated-Image",
        "Image-Size",
    ]

end_experiment()

Ends the experiment. This function is called automatically at the end of __call__ the parameter end_experiment is set to True.

Source code in wandb_addons/diffusers/callbacks/base/base_callback.py 
233
234
235
236
237
238
239
240
241
242
def end_experiment(self):
    """Ends the experiment. This function is called automatically at the end of
    `__call__` the parameter `end_experiment` is set to `True`.
    """
    if self.weave_mode:
        self.stream_table.log(self.table_row)
        self.stream_table.finish()
    elif wandb.run is not None:
        wandb.log({self.table_name: self.wandb_table})
        wandb.finish()

generate(latents) abstractmethod

Generate images from latents. This method must be implemented in the child class.

Source code in wandb_addons/diffusers/callbacks/base/base_callback.py 
152
153
154
155
156
@abstractmethod
def generate(self, latents: torch.FloatTensor) -> List:
    """Generate images from latents. This method must be implemented in the child
    class."""
    pass

initialize_wandb(wandb_project, wandb_entity)

Initializes the wandb run if not already initialized. If weave_mode is True, then a StreamTable is initialized instead of a wandb run. This function is called automatically when the callback is initialized.

Parameters:

Name Type Description Default
wandb_project str

The name of the W&B project.

 required
wandb_entity str

The name of the W&B entity.

 required
Source code in wandb_addons/diffusers/callbacks/base/base_callback.py 
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def initialize_wandb(self, wandb_project: str, wandb_entity: str) -> None:
    """Initializes the wandb run if not already initialized. If `weave_mode` is
    `True`, then a [StreamTable](https://docs.wandb.ai/guides/weave/streamtable) is
    initialized instead of a wandb run. This function is called automatically when
    the callback is initialized.

    Arguments:
        wandb_project (str): The name of the W&B project.
        wandb_entity (str): The name of the W&B entity.
    """
    self.update_configs()
    if self.weave_mode:
        if self.wandb_entity is None:
            wandb.termerror(
                "The parameter wandb_entity must be provided when weave_mode is enabled."
            )
        else:
            self.stream_table = StreamTable(
                f"{self.wandb_entity}/{self.wandb_project}/{self.table_name}"
            )
            self.table_row = []
    else:
        if wandb.run is None:
            if wandb_project is not None:
                wandb.init(
                    project=wandb_project,
                    entity=wandb_entity,
                    job_type=self.job_type,
                    config=self.configs,
                )
            else:
                wandb.termerror("The parameter wandb_project must be provided.")

populate_table_row(prompt, negative_prompt, image)

Populates the table row with the prompt, negative prompt, the generated image, and the configs.

Parameters:

Name Type Description Default
prompt str

The prompt to guide the image generation.

 required
negative_prompt str

The prompt not to guide the image generation.

 required
image Image

The generated image.

 required
Source code in wandb_addons/diffusers/callbacks/base/base_callback.py 
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
def populate_table_row(
    self, prompt: str, negative_prompt: str, image: Image
) -> None:
    """Populates the table row with the prompt, negative prompt, the generated
    image, and the configs.

    Arguments:
        prompt (str): The prompt to guide the image generation.
        negative_prompt (str): The prompt not to guide the image generation.
        image (Image): The generated image.
    """
    width, height = image.size
    if self.weave_mode:
        self.table_row += [
            {
                "Generated-Image": image,
                "Image-Size": {"Width": width, "Height": height},
                "Configs": self.configs,
            }
        ]
    else:
        self.table_row = [
            prompt,
            negative_prompt if negative_prompt is not None else "",
            wandb.Image(image),
            {"Width": width, "Height": height},
        ]

update_configs()

Update the configs as a state of the callback. This function is called inside initialize_wandb().

Source code in wandb_addons/diffusers/callbacks/base/base_callback.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
def update_configs(self) -> None:
    """Update the configs as a state of the callback. This function is called inside
    `initialize_wandb()`.
    """
    pipeline_configs = dict(self.pipeline.config)
    pipeline_configs["scheduler"] = list(pipeline_configs["scheduler"])
    pipeline_configs["scheduler"][1] = dict(self.pipeline.scheduler.config)
    additional_configs = {
        "prompt": self.prompt,
        "negative_prompt": self.negative_prompt
        if self.negative_prompt is not None
        else "",
        "num_inference_steps": self.num_inference_steps,
        "num_images_per_prompt": self.num_images_per_prompt,
        "pipeline": pipeline_configs,
    }
    self.configs = (
        {**self.configs, **additional_configs}
        if self.configs is not None
        else additional_configs
    )

BaseImage2ImageCallback

Bases: BaseDiffusersCallback

Base callback for 🧨 Diffusers logging the results of a DiffusionPipeline for image2image translation task to Weights & Biases.

Parameters:

Name Type Description Default
pipeline DiffusionPipeline

The DiffusionPipeline from diffusers.

 required
prompt Union[str, List[str]]

The prompt or prompts to guide the image generation.

 required
input_images PipelineImageInput

The input image, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between [0, 1] If it's a tensor or a list or tensors, the expected shape should be (B, C, H, W) or (C, H, W). If it is a numpy array or a list of arrays, the expected shape should be (B, H, W, C) or (H, W, C) It can also accept image latents as image, but if passing latents directly it is not encoded again.

 required
wandb_project Optional[str]

The name of the project where you're sending the new run. The project is not necessary to be specified unless the run has automatically been initiatlized before the callback is defined.

 required
wandb_entity Optional[str]

An entity is a username or team name where you're sending runs. This entity must exist before you can send runs there, so make sure to create your account or team in the UI before starting to log runs. If you don't specify an entity, the run will be sent to your default entity, which is usually your username. Change your default entity in your settings under "default location to create new projects".

 None
weave_mode bool

Whether to use log to a weave board instead of W&B dashboard or not. The weave mode logs the configs, generated images and timestamp in a StreamTable instead of a wandb.Table and does not require a wandb run to be initialized in order to start logging. This makes it possible to log muliple generations without having to initialize or terminate runs. Note that the parameter wandb_entity must be explicitly specified in order to use weave mode.

 False
num_inference_steps int

The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference.

 50
num_images_per_prompt Optional[int]

The number of images to generate per prompt.

 1
negative_prompt Optional[Union[str, List[str]]]

The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if guidance_scale is less than 1).

 None
configs Optional[Dict]

Additional configs for the experiment you want to sync, for example, for example, seed could be a good config to be passed here.

 None
Source code in wandb_addons/diffusers/callbacks/base/base_image_to_image_callback.py 
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
class BaseImage2ImageCallback(BaseDiffusersCallback):
    """Base callback for [🧨 Diffusers](https://huggingface.co/docs/diffusers/index)
    logging the results of a
    [`DiffusionPipeline`](https://github.com/huggingface/diffusers/blob/v0.21.0/src/diffusers/pipelines/pipeline_utils.py#L480)
    for image2image translation task to Weights & Biases.

    Arguments:
        pipeline (diffusers.DiffusionPipeline): The `DiffusionPipeline` from
            `diffusers`.
        prompt (Union[str, List[str]]): The prompt or prompts to guide the image
            generation.
        input_images (PipelineImageInput): The input image, numpy array or tensor
            representing an image batch to be used as the starting point. For both numpy
            array and pytorch tensor, the expected value range is between [0, 1] If it's
            a tensor or a list or tensors, the expected shape should be `(B, C, H, W)`
            or `(C, H, W)`. If it is a numpy array or a list of arrays, the expected
            shape should be (B, H, W, C) or (H, W, C) It can also accept image latents as
            image, but if passing latents directly it is not encoded again.
        wandb_project (Optional[str]): The name of the project where you're sending
            the new run. The project is not necessary to be specified unless the run
            has automatically been initiatlized before the callback is defined.
        wandb_entity (Optional[str]): An entity is a username or team name where
            you're sending runs. This entity must exist before you can send runs there,
            so make sure to create your account or team in the UI before starting to
            log runs. If you don't specify an entity, the run will be sent to your
            default entity, which is usually your username. Change your default entity
            in [your settings](https://wandb.ai/settings) under "default location to
            create new projects".
        weave_mode (bool): Whether to use log to a
            [weave board](https://docs.wandb.ai/guides/weave) instead of W&B dashboard
            or not. The weave mode logs the configs, generated images and timestamp in a
            [`StreamTable`](https://docs.wandb.ai/guides/weave/streamtable) instead of a
            `wandb.Table` and does not require a wandb run to be initialized in order to
            start logging. This makes it possible to log muliple generations without
            having to initialize or terminate runs. Note that the parameter
            `wandb_entity` must be explicitly specified in order to use weave mode.
        num_inference_steps (int): The number of denoising steps. More denoising steps
            usually lead to a higher quality image at the expense of slower inference.
        num_images_per_prompt (Optional[int]): The number of images to generate per
            prompt.
        negative_prompt (Optional[Union[str, List[str]]]): The prompt or prompts not
            to guide the image generation. Ignored when not using guidance
            (i.e., ignored if `guidance_scale` is less than `1`).
        configs (Optional[Dict]): Additional configs for the experiment you want to
            sync, for example, for example, `seed` could be a good config to be passed
            here.
    """

    def __init__(
        self,
        pipeline: DiffusionPipeline,
        prompt: Union[str, List[str]],
        input_images: PipelineImageInput,
        wandb_project: str,
        wandb_entity: Optional[str] = None,
        weave_mode: bool = False,
        num_inference_steps: int = 50,
        num_images_per_prompt: Optional[int] = 1,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        configs: Optional[Dict] = None,
        **kwargs,
    ) -> None:
        self.job_type = "image-to-image"
        super().__init__(
            pipeline=pipeline,
            prompt=prompt,
            wandb_project=wandb_project,
            wandb_entity=wandb_entity,
            weave_mode=weave_mode,
            num_inference_steps=num_inference_steps,
            num_images_per_prompt=num_images_per_prompt,
            negative_prompt=negative_prompt,
            configs=configs,
            **kwargs,
        )
        self.input_images = self.postprocess_input_images(input_images)

    def initialize_wandb(self, wandb_project, wandb_entity) -> None:
        """Initializes the wandb run if not already initialized. If `weave_mode` is
        `True`, then a [StreamTable](https://docs.wandb.ai/guides/weave/streamtable) is
        initialized instead of a wandb run. This function is called automatically when
        the callback is initialized.

        Arguments:
            wandb_project (str): The name of the W&B project.
            wandb_entity (str): The name of the W&B entity.
        """
        self.job_type = "image-to-image"
        self.table_name = "Image-To-Image"
        super().initialize_wandb(wandb_project, wandb_entity)

    def postprocess_input_images(
        self, input_images: Union[torch.Tensor, Image.Image, np.array]
    ) -> Image.Image:
        """Postprocess input images to be logged to the W&B Table/StreamTable.

        Arguments:
            input_images (Union[torch.Tensor, Image.Image, np.array]): The input images
                to be postprocessed.
        """
        if isinstance(input_images, torch.Tensor):
            input_images = self.pipeline.image_processor.pt_to_numpy(input_images)
            input_images = self.pipeline.image_processor.numpy_to_pil(input_images)
        elif isinstance(input_images, Image.Image):
            input_images = [input_images]
        elif isinstance(input_images, np.array):
            input_images = self.pipeline.image_processor.numpy_to_pil(input_images)
        return input_images

    def build_wandb_table(self) -> None:
        """Specifies the columns of the wandb table if not in weave mode. This function
        is called automatically when the callback is initialized.
        """
        super().build_wandb_table()
        self.table_columns = ["Input-Image", "Input-Image-Size"] + self.table_columns

    def populate_table_row(
        self, input_image: Image.Image, prompt: str, negative_prompt: str, image: Any
    ) -> None:
        """Populates the table row with the input image, prompt, negative prompt, the
        generated image, and the configs.

        Arguments:
            input_image (Image): The input image.s
            prompt (str): The prompt to guide the image generation.
            negative_prompt (str): The prompt not to guide the image generation.
            image (Image): The generated image.
        """
        input_width, input_height = input_image.size
        generated_width, generated_height = image.size
        self.table_row = (
            {
                "Input-Image": input_image,
                "Input-Image-Size": {"Width": input_width, "Height": input_height},
                "Prompt": prompt,
                "Negative-Prompt": negative_prompt
                if negative_prompt is not None
                else "",
                "Generated-Image": image,
                "Generated-Image-Size": {
                    "Width": generated_width,
                    "Height": generated_height,
                },
                "Configs": self.configs,
            }
            if self.weave_mode
            else [
                wandb.Image(input_image),
                {"Width": input_width, "Height": input_height},
                prompt,
                negative_prompt if negative_prompt is not None else "",
                wandb.Image(image),
                {"Width": generated_width, "Height": generated_height},
            ]
        )

    def __call__(
        self,
        step: int,
        timestep: int,
        latents: torch.FloatTensor,
        end_experiment: bool = True,
    ):
        """A function that will be called every `callback_steps` steps during
        inference with the `diffusers.DiffusionPipeline`.

        Arguments:
            step (int): The current step of the inference.
            timestep (int): The current timestep of the inference.
            latents (torch.FloatTensor): The latent tensor used to generate the image.
            end_experiment (bool): Whether to end the experiment automatically or not
                after the pipeline is executed.
        """
        if step == self.starting_step:
            self.at_initial_step()
        if step == self.log_step:
            images = self.generate(latents)
            prompt_logging = (
                self.prompt if isinstance(self.prompt, list) else [self.prompt]
            )
            negative_prompt_logging = (
                self.negative_prompt
                if isinstance(self.negative_prompt, list)
                else [self.negative_prompt] * len(prompt_logging)
            )
            images = chunkify(images, len(prompt_logging))
            for idx in range(len(prompt_logging)):
                for image in images[idx]:
                    self.populate_table_row(
                        self.input_images[0],
                        prompt_logging[idx],
                        negative_prompt_logging[idx],
                        image,
                    )
                    if self.weave_mode:
                        self.stream_table.log(self.table_row)
                    else:
                        self.wandb_table.add_data(*self.table_row)
                        wandb.log({"Generated-Images": wandb.Image(image)})
            if end_experiment:
                self.end_experiment()

__call__(step, timestep, latents, end_experiment=True)

A function that will be called every callback_steps steps during inference with the diffusers.DiffusionPipeline.

Parameters:

Name Type Description Default
step int

The current step of the inference.

 required
timestep int

The current timestep of the inference.

 required
latents FloatTensor

The latent tensor used to generate the image.

 required
end_experiment bool

Whether to end the experiment automatically or not after the pipeline is executed.

 True
Source code in wandb_addons/diffusers/callbacks/base/base_image_to_image_callback.py 
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
def __call__(
    self,
    step: int,
    timestep: int,
    latents: torch.FloatTensor,
    end_experiment: bool = True,
):
    """A function that will be called every `callback_steps` steps during
    inference with the `diffusers.DiffusionPipeline`.

    Arguments:
        step (int): The current step of the inference.
        timestep (int): The current timestep of the inference.
        latents (torch.FloatTensor): The latent tensor used to generate the image.
        end_experiment (bool): Whether to end the experiment automatically or not
            after the pipeline is executed.
    """
    if step == self.starting_step:
        self.at_initial_step()
    if step == self.log_step:
        images = self.generate(latents)
        prompt_logging = (
            self.prompt if isinstance(self.prompt, list) else [self.prompt]
        )
        negative_prompt_logging = (
            self.negative_prompt
            if isinstance(self.negative_prompt, list)
            else [self.negative_prompt] * len(prompt_logging)
        )
        images = chunkify(images, len(prompt_logging))
        for idx in range(len(prompt_logging)):
            for image in images[idx]:
                self.populate_table_row(
                    self.input_images[0],
                    prompt_logging[idx],
                    negative_prompt_logging[idx],
                    image,
                )
                if self.weave_mode:
                    self.stream_table.log(self.table_row)
                else:
                    self.wandb_table.add_data(*self.table_row)
                    wandb.log({"Generated-Images": wandb.Image(image)})
        if end_experiment:
            self.end_experiment()

build_wandb_table()

Specifies the columns of the wandb table if not in weave mode. This function is called automatically when the callback is initialized.

Source code in wandb_addons/diffusers/callbacks/base/base_image_to_image_callback.py 
123
124
125
126
127
128
def build_wandb_table(self) -> None:
    """Specifies the columns of the wandb table if not in weave mode. This function
    is called automatically when the callback is initialized.
    """
    super().build_wandb_table()
    self.table_columns = ["Input-Image", "Input-Image-Size"] + self.table_columns

initialize_wandb(wandb_project, wandb_entity)

Initializes the wandb run if not already initialized. If weave_mode is True, then a StreamTable is initialized instead of a wandb run. This function is called automatically when the callback is initialized.

Parameters:

Name Type Description Default
wandb_project str

The name of the W&B project.

 required
wandb_entity str

The name of the W&B entity.

 required
Source code in wandb_addons/diffusers/callbacks/base/base_image_to_image_callback.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def initialize_wandb(self, wandb_project, wandb_entity) -> None:
    """Initializes the wandb run if not already initialized. If `weave_mode` is
    `True`, then a [StreamTable](https://docs.wandb.ai/guides/weave/streamtable) is
    initialized instead of a wandb run. This function is called automatically when
    the callback is initialized.

    Arguments:
        wandb_project (str): The name of the W&B project.
        wandb_entity (str): The name of the W&B entity.
    """
    self.job_type = "image-to-image"
    self.table_name = "Image-To-Image"
    super().initialize_wandb(wandb_project, wandb_entity)

populate_table_row(input_image, prompt, negative_prompt, image)

Populates the table row with the input image, prompt, negative prompt, the generated image, and the configs.

Parameters:

Name Type Description Default
input_image Image

The input image.s

 required
prompt str

The prompt to guide the image generation.

 required
negative_prompt str

The prompt not to guide the image generation.

 required
image Image

The generated image.

 required
Source code in wandb_addons/diffusers/callbacks/base/base_image_to_image_callback.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
def populate_table_row(
    self, input_image: Image.Image, prompt: str, negative_prompt: str, image: Any
) -> None:
    """Populates the table row with the input image, prompt, negative prompt, the
    generated image, and the configs.

    Arguments:
        input_image (Image): The input image.s
        prompt (str): The prompt to guide the image generation.
        negative_prompt (str): The prompt not to guide the image generation.
        image (Image): The generated image.
    """
    input_width, input_height = input_image.size
    generated_width, generated_height = image.size
    self.table_row = (
        {
            "Input-Image": input_image,
            "Input-Image-Size": {"Width": input_width, "Height": input_height},
            "Prompt": prompt,
            "Negative-Prompt": negative_prompt
            if negative_prompt is not None
            else "",
            "Generated-Image": image,
            "Generated-Image-Size": {
                "Width": generated_width,
                "Height": generated_height,
            },
            "Configs": self.configs,
        }
        if self.weave_mode
        else [
            wandb.Image(input_image),
            {"Width": input_width, "Height": input_height},
            prompt,
            negative_prompt if negative_prompt is not None else "",
            wandb.Image(image),
            {"Width": generated_width, "Height": generated_height},
        ]
    )

postprocess_input_images(input_images)

Postprocess input images to be logged to the W&B Table/StreamTable.

Parameters:

Name Type Description Default
input_images Union[Tensor, Image, array]

The input images to be postprocessed.

 required
Source code in wandb_addons/diffusers/callbacks/base/base_image_to_image_callback.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
def postprocess_input_images(
    self, input_images: Union[torch.Tensor, Image.Image, np.array]
) -> Image.Image:
    """Postprocess input images to be logged to the W&B Table/StreamTable.

    Arguments:
        input_images (Union[torch.Tensor, Image.Image, np.array]): The input images
            to be postprocessed.
    """
    if isinstance(input_images, torch.Tensor):
        input_images = self.pipeline.image_processor.pt_to_numpy(input_images)
        input_images = self.pipeline.image_processor.numpy_to_pil(input_images)
    elif isinstance(input_images, Image.Image):
        input_images = [input_images]
    elif isinstance(input_images, np.array):
        input_images = self.pipeline.image_processor.numpy_to_pil(input_images)
    return input_images

BaseMultiPipelineCallback

Bases: BaseDiffusersCallback

Base callback for 🧨 Diffusers logging the results of multiple DiffusionPipelines to Weights & Biases.

Parameters:

Name Type Description Default
pipeline DiffusionPipeline

The DiffusionPipeline from diffusers.

 required
prompt Union[str, List[str]]

The prompt or prompts to guide the image generation.

 required
wandb_project Optional[str]

The name of the project where you're sending the new run. The project is not necessary to be specified unless the run has automatically been initiatlized before the callback is defined.

 required
wandb_entity Optional[str]

An entity is a username or team name where you're sending runs. This entity must exist before you can send runs there, so make sure to create your account or team in the UI before starting to log runs. If you don't specify an entity, the run will be sent to your default entity, which is usually your username. Change your default entity in your settings under "default location to create new projects".

 None
weave_mode bool

Whether to use log to a weave board instead of W&B dashboard or not. The weave mode logs the configs, generated images and timestamp in a StreamTable instead of a wandb.Table and does not require a wandb run to be initialized in order to start logging. This makes it possible to log muliple generations without having to initialize or terminate runs. Note that the parameter wandb_entity must be explicitly specified in order to use weave mode.

 False
num_inference_steps int

The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference.

 50
num_images_per_prompt Optional[int]

The number of images to generate per prompt.

 1
negative_prompt Optional[Union[str, List[str]]]

The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if guidance_scale is less than 1).

 None
initial_stage_name Optional[str]

The name of the initial stage. If not specified, it would be set to "stage_1".

 None
configs Optional[Dict]

Additional configs for the experiment you want to sync, for example, for example, seed could be a good config to be passed here.

 None
Source code in wandb_addons/diffusers/callbacks/base/base_multi_pipeline_callback.py 
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
class BaseMultiPipelineCallback(BaseDiffusersCallback):
    """Base callback for [🧨 Diffusers](https://huggingface.co/docs/diffusers/index)
    logging the results of multiple
    [`DiffusionPipeline`](https://github.com/huggingface/diffusers/blob/v0.21.0/src/diffusers/pipelines/pipeline_utils.py#L480)s
    to Weights & Biases.

    Arguments:
        pipeline (diffusers.DiffusionPipeline): The `DiffusionPipeline` from
            `diffusers`.
        prompt (Union[str, List[str]]): The prompt or prompts to guide the image
            generation.
        wandb_project (Optional[str]): The name of the project where you're sending
            the new run. The project is not necessary to be specified unless the run
            has automatically been initiatlized before the callback is defined.
        wandb_entity (Optional[str]): An entity is a username or team name where
            you're sending runs. This entity must exist before you can send runs there,
            so make sure to create your account or team in the UI before starting to
            log runs. If you don't specify an entity, the run will be sent to your
            default entity, which is usually your username. Change your default entity
            in [your settings](https://wandb.ai/settings) under "default location to
            create new projects".
        weave_mode (bool): Whether to use log to a
            [weave board](https://docs.wandb.ai/guides/weave) instead of W&B dashboard
            or not. The weave mode logs the configs, generated images and timestamp in a
            [`StreamTable`](https://docs.wandb.ai/guides/weave/streamtable) instead of a
            `wandb.Table` and does not require a wandb run to be initialized in order to
            start logging. This makes it possible to log muliple generations without
            having to initialize or terminate runs. Note that the parameter
            `wandb_entity` must be explicitly specified in order to use weave mode.
        num_inference_steps (int): The number of denoising steps. More denoising steps
            usually lead to a higher quality image at the expense of slower inference.
        num_images_per_prompt (Optional[int]): The number of images to generate per
            prompt.
        negative_prompt (Optional[Union[str, List[str]]]): The prompt or prompts not
            to guide the image generation. Ignored when not using guidance
            (i.e., ignored if `guidance_scale` is less than `1`).
        initial_stage_name (Optional[str]): The name of the initial stage. If not
            specified, it would be set to `"stage_1"`.
        configs (Optional[Dict]): Additional configs for the experiment you want to
            sync, for example, for example, `seed` could be a good config to be passed
            here.
    """

    def __init__(
        self,
        pipeline: DiffusionPipeline,
        prompt: Union[str, List[str]],
        wandb_project: str,
        wandb_entity: Optional[str] = None,
        weave_mode: bool = False,
        num_inference_steps: int = 50,
        num_images_per_prompt: Optional[int] = 1,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        initial_stage_name: Optional[str] = None,
        configs: Optional[Dict] = None,
        **kwargs,
    ) -> None:
        self.stage_name = (
            initial_stage_name if initial_stage_name is not None else "stage_1"
        )
        self.stage_counter = 1
        self.original_configs = {} if configs is None else configs
        super().__init__(
            pipeline,
            prompt,
            wandb_project,
            wandb_entity,
            weave_mode,
            num_inference_steps,
            num_images_per_prompt,
            negative_prompt,
            configs,
            **kwargs,
        )
        self.table_row = {}

    def update_configs(self) -> None:
        """Update the configs as a state of the callback. This function is called inside
        `initialize_wandb()`. In this function, the configs regarding the base pipeline
        are updated as well.
        """
        pipeline_configs = dict(self.pipeline.config)
        pipeline_configs["scheduler"] = list(pipeline_configs["scheduler"])
        pipeline_configs["scheduler"][1] = dict(self.pipeline.scheduler.config)
        additional_configs = {
            self.stage_name: {
                "pipeline": pipeline_configs,
                "num_inference_steps": self.num_inference_steps,
                "prompt": self.prompt,
                "negative_prompt": self.negative_prompt
                if self.negative_prompt is not None
                else "",
                "num_images_per_prompt": self.num_images_per_prompt,
                "stage-name": self.stage_name,
                "stage-sequence": self.stage_counter,
                **self.original_configs,
            },
        }
        self.configs = additional_configs

    def add_stage(
        self,
        pipeline: DiffusionPipeline,
        num_inference_steps: Optional[int] = None,
        stage_name: Optional[str] = None,
        configs: Optional[Dict] = None,
    ) -> None:
        """Add a new stage to the callback to log the results of a new pipeline in a
        multi-pipeline workflow.

        Arguments:
            pipeline (diffusers.DiffusionPipeline): The `DiffusionPipeline` from
                for the new stage.
            num_inference_steps (Optional[int]): The number of denoising steps for the
                new stage. More denoising steps usually lead to a higher quality image
                at the expense of slower inference.
            stage_name (Optional[str]): The name of the new stage. If not specified,
                it would be set to `"stage_{stage_counter}"`.
            configs (Optional[Dict]): Additional configs for the new stage you want to
                sync, for example, for example, `seed` could be a good config to be
                passed here.
        """
        self.pipeline = pipeline
        self.num_inference_steps = (
            num_inference_steps
            if num_inference_steps is not None
            else self.num_inference_steps
        )
        self.stage_counter += 1
        self.stage_name = (
            stage_name if stage_name is not None else f"stage_{self.stage_counter}"
        )
        pipeline_configs = dict(self.pipeline.config)
        pipeline_configs["scheduler"] = list(pipeline_configs["scheduler"])
        pipeline_configs["scheduler"][1] = dict(self.pipeline.scheduler.config)
        additional_configs = {
            self.stage_name: {
                "pipeline": pipeline_configs,
                "num_inference_steps": self.num_inference_steps,
                "prompt": self.prompt,
                "negative_prompt": self.negative_prompt
                if self.negative_prompt is not None
                else "",
                "num_images_per_prompt": self.num_images_per_prompt,
                "stage-name": self.stage_name,
                "stage-sequence": self.stage_counter,
                **self.original_configs,
            },
        }
        if configs is not None:
            additional_configs[self.stage_name].update(configs)
        self.configs.update(additional_configs)
        if wandb.run is not None:
            wandb.config.update(additional_configs)

    def at_initial_step(self):
        """A function that will be called at the initial step of the denoising loop
        during inference."""
        if self.stage_counter == 1:
            super().at_initial_step()

    def build_wandb_table(self) -> None:
        """Specifies the columns of the wandb table if not in weave mode. This function
        is called automatically when the callback is initialized.
        """
        super().build_wandb_table()
        self.table_columns = ["Stage-Sequence", "Stage-Name"] + self.table_columns

    def populate_table_row(self, prompt: str, negative_prompt: str, image: Any) -> None:
        """Populates the table row with the prompt, negative prompt, the generated
        image, and the configs.

        Arguments:
            prompt (str): The prompt to guide the image generation.
            negative_prompt (str): The prompt not to guide the image generation.
            image (Image): The generated image.
        """
        width, height = image.size
        if self.weave_mode:
            self.table_row.update(
                {
                    self.stage_name: {
                        "Generated-Image": image,
                        "Image-Size": {"Width": width, "Height": height},
                        "Configs": self.configs[self.stage_name],
                    }
                }
            )
        else:
            self.table_row = [
                self.stage_counter,
                self.stage_name,
                prompt,
                negative_prompt if negative_prompt is not None else "",
                wandb.Image(image),
                {"Width": width, "Height": height},
            ]

    def end_experiment(self):
        """Ends the experiment. This function is called automatically at the end of
        `__call__` the parameter `end_experiment` is set to `True`.
        """
        if self.weave_mode:
            self.table_row = {"Experiment": self.table_row}
            self.stream_table.log(self.table_row)
            self.stream_table.finish()
        elif wandb.run is not None:
            wandb.log({self.table_name: self.wandb_table})
            wandb.finish()

add_stage(pipeline, num_inference_steps=None, stage_name=None, configs=None)

Add a new stage to the callback to log the results of a new pipeline in a multi-pipeline workflow.

Parameters:

Name Type Description Default
pipeline DiffusionPipeline

The DiffusionPipeline from for the new stage.

 required
num_inference_steps Optional[int]

The number of denoising steps for the new stage. More denoising steps usually lead to a higher quality image at the expense of slower inference.

 None
stage_name Optional[str]

The name of the new stage. If not specified, it would be set to "stage_{stage_counter}".

 None
configs Optional[Dict]

Additional configs for the new stage you want to sync, for example, for example, seed could be a good config to be passed here.

 None
Source code in wandb_addons/diffusers/callbacks/base/base_multi_pipeline_callback.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def add_stage(
    self,
    pipeline: DiffusionPipeline,
    num_inference_steps: Optional[int] = None,
    stage_name: Optional[str] = None,
    configs: Optional[Dict] = None,
) -> None:
    """Add a new stage to the callback to log the results of a new pipeline in a
    multi-pipeline workflow.

    Arguments:
        pipeline (diffusers.DiffusionPipeline): The `DiffusionPipeline` from
            for the new stage.
        num_inference_steps (Optional[int]): The number of denoising steps for the
            new stage. More denoising steps usually lead to a higher quality image
            at the expense of slower inference.
        stage_name (Optional[str]): The name of the new stage. If not specified,
            it would be set to `"stage_{stage_counter}"`.
        configs (Optional[Dict]): Additional configs for the new stage you want to
            sync, for example, for example, `seed` could be a good config to be
            passed here.
    """
    self.pipeline = pipeline
    self.num_inference_steps = (
        num_inference_steps
        if num_inference_steps is not None
        else self.num_inference_steps
    )
    self.stage_counter += 1
    self.stage_name = (
        stage_name if stage_name is not None else f"stage_{self.stage_counter}"
    )
    pipeline_configs = dict(self.pipeline.config)
    pipeline_configs["scheduler"] = list(pipeline_configs["scheduler"])
    pipeline_configs["scheduler"][1] = dict(self.pipeline.scheduler.config)
    additional_configs = {
        self.stage_name: {
            "pipeline": pipeline_configs,
            "num_inference_steps": self.num_inference_steps,
            "prompt": self.prompt,
            "negative_prompt": self.negative_prompt
            if self.negative_prompt is not None
            else "",
            "num_images_per_prompt": self.num_images_per_prompt,
            "stage-name": self.stage_name,
            "stage-sequence": self.stage_counter,
            **self.original_configs,
        },
    }
    if configs is not None:
        additional_configs[self.stage_name].update(configs)
    self.configs.update(additional_configs)
    if wandb.run is not None:
        wandb.config.update(additional_configs)

at_initial_step()

A function that will be called at the initial step of the denoising loop during inference.

Source code in wandb_addons/diffusers/callbacks/base/base_multi_pipeline_callback.py 
164
165
166
167
168
def at_initial_step(self):
    """A function that will be called at the initial step of the denoising loop
    during inference."""
    if self.stage_counter == 1:
        super().at_initial_step()

build_wandb_table()

Specifies the columns of the wandb table if not in weave mode. This function is called automatically when the callback is initialized.

Source code in wandb_addons/diffusers/callbacks/base/base_multi_pipeline_callback.py 
170
171
172
173
174
175
def build_wandb_table(self) -> None:
    """Specifies the columns of the wandb table if not in weave mode. This function
    is called automatically when the callback is initialized.
    """
    super().build_wandb_table()
    self.table_columns = ["Stage-Sequence", "Stage-Name"] + self.table_columns

end_experiment()

Ends the experiment. This function is called automatically at the end of __call__ the parameter end_experiment is set to True.

Source code in wandb_addons/diffusers/callbacks/base/base_multi_pipeline_callback.py 
207
208
209
210
211
212
213
214
215
216
217
def end_experiment(self):
    """Ends the experiment. This function is called automatically at the end of
    `__call__` the parameter `end_experiment` is set to `True`.
    """
    if self.weave_mode:
        self.table_row = {"Experiment": self.table_row}
        self.stream_table.log(self.table_row)
        self.stream_table.finish()
    elif wandb.run is not None:
        wandb.log({self.table_name: self.wandb_table})
        wandb.finish()

populate_table_row(prompt, negative_prompt, image)

Populates the table row with the prompt, negative prompt, the generated image, and the configs.

Parameters:

Name Type Description Default
prompt str

The prompt to guide the image generation.

 required
negative_prompt str

The prompt not to guide the image generation.

 required
image Image

The generated image.

 required
Source code in wandb_addons/diffusers/callbacks/base/base_multi_pipeline_callback.py 
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
def populate_table_row(self, prompt: str, negative_prompt: str, image: Any) -> None:
    """Populates the table row with the prompt, negative prompt, the generated
    image, and the configs.

    Arguments:
        prompt (str): The prompt to guide the image generation.
        negative_prompt (str): The prompt not to guide the image generation.
        image (Image): The generated image.
    """
    width, height = image.size
    if self.weave_mode:
        self.table_row.update(
            {
                self.stage_name: {
                    "Generated-Image": image,
                    "Image-Size": {"Width": width, "Height": height},
                    "Configs": self.configs[self.stage_name],
                }
            }
        )
    else:
        self.table_row = [
            self.stage_counter,
            self.stage_name,
            prompt,
            negative_prompt if negative_prompt is not None else "",
            wandb.Image(image),
            {"Width": width, "Height": height},
        ]

update_configs()

Update the configs as a state of the callback. This function is called inside initialize_wandb(). In this function, the configs regarding the base pipeline are updated as well.

Source code in wandb_addons/diffusers/callbacks/base/base_multi_pipeline_callback.py 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
def update_configs(self) -> None:
    """Update the configs as a state of the callback. This function is called inside
    `initialize_wandb()`. In this function, the configs regarding the base pipeline
    are updated as well.
    """
    pipeline_configs = dict(self.pipeline.config)
    pipeline_configs["scheduler"] = list(pipeline_configs["scheduler"])
    pipeline_configs["scheduler"][1] = dict(self.pipeline.scheduler.config)
    additional_configs = {
        self.stage_name: {
            "pipeline": pipeline_configs,
            "num_inference_steps": self.num_inference_steps,
            "prompt": self.prompt,
            "negative_prompt": self.negative_prompt
            if self.negative_prompt is not None
            else "",
            "num_images_per_prompt": self.num_images_per_prompt,
            "stage-name": self.stage_name,
            "stage-sequence": self.stage_counter,
            **self.original_configs,
        },
    }
    self.configs = additional_configs