Init demo.

app.py

@@ -1,7 +1,39 @@
-import gradio as gr
-
-def greet(name):
-    return "Hello " + name + "!!"
-
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+import gradio as gr
+from webui import (
+    create_interface_texture_synthesis,
+    create_interface_style_t2i,
+    create_interface_style_transfer,
+    Runner
+)
+
+
+import os
+os.environ["no_proxy"] = "localhost,127.0.0.1,::1"
+
+
+def main():
+    runner = Runner()
+
+    with gr.Blocks(analytics_enabled=False,
+                   title='Attention Distillation',
+                   ) as demo:
+
+        md_txt = "# Attention Distillation" \
+                 "\nOfficial demo of the paper [Attention Distillation: A Unified Approach to Visual Characteristics Transfer](https://arxiv.org/abs/2502.20235)"
+        gr.Markdown(md_txt)
+
+        with gr.Tabs(selected='tab_style_transfer'):
+            with gr.TabItem("Style Transfer", id='tab_style_transfer'):
+                create_interface_style_transfer(runner=runner)
+
+            with gr.TabItem("Style-Specific Text-to-Image Generation", id='tab_style_t2i'):
+                create_interface_style_t2i(runner=runner)
+
+            with gr.TabItem("Texture Synthesis", id='tab_texture_syn'):
+                create_interface_texture_synthesis(runner=runner)
+            
+        demo.launch(share=False, debug=False)
+
+
+if __name__ == '__main__':
+    main()

losses.py

@@ -0,0 +1,56 @@
+import math
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+loss_fn = torch.nn.L1Loss()
+
+
+def ad_loss(
+    q_list, ks_list, vs_list, self_out_list, scale=1, source_mask=None, target_mask=None
+):
+    loss = 0
+    attn_mask = None
+    for q, ks, vs, self_out in zip(q_list, ks_list, vs_list, self_out_list):
+        if source_mask is not None and target_mask is not None:
+            w = h = int(np.sqrt(q.shape[2]))
+            mask_1 = torch.flatten(F.interpolate(source_mask, size=(h, w)))
+            mask_2 = torch.flatten(F.interpolate(target_mask, size=(h, w)))
+            attn_mask = mask_1.unsqueeze(0) == mask_2.unsqueeze(1)
+            attn_mask=attn_mask.to(q.device)
+
+        target_out = F.scaled_dot_product_attention(
+            q * scale,
+            torch.cat(torch.chunk(ks, ks.shape[0]), 2).repeat(q.shape[0], 1, 1, 1),
+            torch.cat(torch.chunk(vs, vs.shape[0]), 2).repeat(q.shape[0], 1, 1, 1),
+            attn_mask=attn_mask
+        )
+        loss += loss_fn(self_out, target_out.detach())
+    return loss
+
+
+
+def q_loss(q_list, qc_list):
+    loss = 0
+    for q, qc in zip(q_list, qc_list):
+        loss += loss_fn(q, qc.detach())
+    return loss
+
+# weight = 200
+def qk_loss(q_list, k_list, qc_list, kc_list):
+    loss = 0
+    for q, k, qc, kc in zip(q_list, k_list, qc_list, kc_list):
+        scale_factor = 1 / math.sqrt(q.size(-1))
+        self_map = torch.softmax(q @ k.transpose(-2, -1) * scale_factor, dim=-1)
+        target_map = torch.softmax(qc @ kc.transpose(-2, -1) * scale_factor, dim=-1)
+        loss += loss_fn(self_map, target_map.detach())
+    return loss
+
+# weight = 1
+def qkv_loss(q_list, k_list, vc_list, c_out_list):
+    loss = 0
+    for q, k, vc, target_out in zip(q_list, k_list, vc_list, c_out_list):
+        self_out = F.scaled_dot_product_attention(q, k, vc)
+        loss += loss_fn(self_out, target_out.detach())
+    return loss

pipeline_sd.py

@@ -0,0 +1,680 @@
+import copy
+import math
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import utils
+from accelerate import Accelerator
+from diffusers import StableDiffusionPipeline
+from diffusers.image_processor import PipelineImageInput
+from losses import *
+from tqdm import tqdm
+
+
+class ADPipeline(StableDiffusionPipeline):
+    def freeze(self):
+        self.vae.requires_grad_(False)
+        self.unet.requires_grad_(False)
+        self.text_encoder.requires_grad_(False)
+        self.classifier.requires_grad_(False)
+
+    @torch.no_grad()
+    def image2latent(self, image):
+        dtype = next(self.vae.parameters()).dtype
+        device = self._execution_device
+        image = image.to(device=device, dtype=dtype) * 2.0 - 1.0
+        latent = self.vae.encode(image)["latent_dist"].mean
+        latent = latent * self.vae.config.scaling_factor
+        return latent
+
+    @torch.no_grad()
+    def latent2image(self, latent):
+        dtype = next(self.vae.parameters()).dtype
+        device = self._execution_device
+        latent = latent.to(device=device, dtype=dtype)
+        latent = latent / self.vae.config.scaling_factor
+        image = self.vae.decode(latent)[0]
+        return (image * 0.5 + 0.5).clamp(0, 1)
+
+    def init(self, enable_gradient_checkpoint):
+        self.freeze()
+        weight_dtype = torch.float32
+        if self.accelerator.mixed_precision == "fp16":
+            weight_dtype = torch.float16
+        elif self.accelerator.mixed_precision == "bf16":
+            weight_dtype = torch.bfloat16
+
+        # Move unet, vae and text_encoder to device and cast to weight_dtype
+        self.unet.to(self.accelerator.device, dtype=weight_dtype)
+        self.vae.to(self.accelerator.device, dtype=weight_dtype)
+        self.text_encoder.to(self.accelerator.device, dtype=weight_dtype)
+        self.classifier.to(self.accelerator.device, dtype=weight_dtype)
+        self.classifier = self.accelerator.prepare(self.classifier)
+        if enable_gradient_checkpoint:
+            self.classifier.enable_gradient_checkpointing()
+
+    def sample(
+        self,
+        lr=0.05,
+        iters=1,
+        attn_scale=1,
+        adain=False,
+        weight=0.25,
+        controller=None,
+        style_image=None,
+        content_image=None,
+        mixed_precision="no",
+        start_time=999,
+        enable_gradient_checkpoint=False,
+        prompt: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        guidance_rescale: float = 0.0,
+        clip_skip: Optional[int] = None,
+        **kwargs,
+    ):
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+        self._guidance_scale = guidance_scale
+        self._guidance_rescale = guidance_rescale
+        self._clip_skip = clip_skip
+        self._cross_attention_kwargs = cross_attention_kwargs
+        self._interrupt = False
+
+        self.accelerator = Accelerator(
+            mixed_precision=mixed_precision, gradient_accumulation_steps=1
+        )
+        self.init(enable_gradient_checkpoint)
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        # 3. Encode input prompt
+        lora_scale = (
+            self.cross_attention_kwargs.get("scale", None)
+            if self.cross_attention_kwargs is not None
+            else None
+        )
+        do_cfg = guidance_scale > 1.0
+
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            do_cfg,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            lora_scale=lora_scale,
+            clip_skip=self.clip_skip,
+        )
+
+        # For classifier free guidance, we need to do two forward passes.
+        # Here we concatenate the unconditional and text embeddings into a single batch
+        # to avoid doing two forward passes
+        if do_cfg:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+                do_cfg,
+            )
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6.1 Add image embeds for IP-Adapter
+        added_cond_kwargs = (
+            {"image_embeds": image_embeds}
+            if (ip_adapter_image is not None or ip_adapter_image_embeds is not None)
+            else None
+        )
+
+        # 6.2 Optionally get Guidance Scale Embedding
+        timestep_cond = None
+        if self.unet.config.time_cond_proj_dim is not None:
+            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(
+                batch_size * num_images_per_prompt
+            )
+            timestep_cond = self.get_guidance_scale_embedding(
+                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
+            ).to(device=device, dtype=latents.dtype)
+
+        self.scheduler.set_timesteps(num_inference_steps)
+        timesteps = self.scheduler.timesteps
+        self.style_latent = self.image2latent(style_image)
+        if content_image is not None:
+            self.content_latent = self.image2latent(content_image)
+        else:
+            self.content_latent = None
+        null_embeds = self.encode_prompt("", device, 1, False)[0]
+        self.null_embeds = null_embeds
+        self.null_embeds_for_latents = torch.cat([null_embeds] * latents.shape[0])
+        self.null_embeds_for_style = torch.cat(
+            [null_embeds] * self.style_latent.shape[0]
+        )
+        
+        self.adain = adain
+        self.attn_scale = attn_scale
+        self.cache = utils.DataCache()
+        self.controller = controller
+        utils.register_attn_control(
+            self.classifier, controller=self.controller, cache=self.cache
+        )
+        print("Total self attention layers of Unet: ", controller.num_self_layers)
+        print("Self attention layers for AD: ", controller.self_layers)
+
+        pbar = tqdm(timesteps, desc="Sample")
+        for i, t in enumerate(pbar):
+            with torch.no_grad():
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if do_cfg else latents
+                latent_model_input = self.scheduler.scale_model_input(
+                    latent_model_input, t
+                )
+                # predict the noise residual
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    timestep_cond=timestep_cond,
+                    cross_attention_kwargs=self.cross_attention_kwargs,
+                    added_cond_kwargs=added_cond_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # perform guidance
+                if do_cfg:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (
+                        noise_pred_text - noise_pred_uncond
+                    )
+                latents = self.scheduler.step(
+                    noise_pred, t, latents, return_dict=False
+                )[0]
+            if iters > 0 and t < start_time:
+                latents = self.AD(latents, t, lr, iters, pbar, weight)
+                
+        images = self.latent2image(latents)
+        # Offload all models
+        self.maybe_free_model_hooks()
+        return images
+
+    def optimize(
+        self,
+        latents=None,
+        attn_scale=1.0,
+        lr=0.05,
+        iters=1,
+        weight=0,
+        width=512,
+        height=512,
+        batch_size=1,
+        controller=None,
+        style_image=None,
+        content_image=None,
+        mixed_precision="no",
+        num_inference_steps=50,
+        enable_gradient_checkpoint=False,
+        source_mask=None,
+        target_mask=None,
+    ):
+        height = height // self.vae_scale_factor
+        width = width // self.vae_scale_factor
+
+        self.accelerator = Accelerator(
+            mixed_precision=mixed_precision, gradient_accumulation_steps=1
+        )
+        self.init(enable_gradient_checkpoint)
+
+        style_latent = self.image2latent(style_image)
+        latents = torch.randn((batch_size, 4, height, width), device=self.device)
+        null_embeds = self.encode_prompt("", self.device, 1, False)[0]
+        null_embeds_for_latents = null_embeds.repeat(latents.shape[0], 1, 1)
+        null_embeds_for_style = null_embeds.repeat(style_latent.shape[0], 1, 1)
+
+        if content_image is not None:
+            content_latent = self.image2latent(content_image)
+            latents = torch.cat([content_latent.clone()] * batch_size)
+            null_embeds_for_content = null_embeds.repeat(content_latent.shape[0], 1, 1)
+
+        self.cache = utils.DataCache()
+        self.controller = controller
+        utils.register_attn_control(
+            self.classifier, controller=self.controller, cache=self.cache
+        )
+        print("Total self attention layers of Unet: ", controller.num_self_layers)
+        print("Self attention layers for AD: ", controller.self_layers)
+
+        self.scheduler.set_timesteps(num_inference_steps)
+        timesteps = self.scheduler.timesteps
+        latents = latents.detach().float()
+        optimizer = torch.optim.Adam([latents.requires_grad_()], lr=lr)
+        optimizer = self.accelerator.prepare(optimizer)
+        pbar = tqdm(timesteps, desc="Optimize")
+        for i, t in enumerate(pbar):
+            # t = torch.tensor([1], device=self.device)
+            with torch.no_grad():
+                qs_list, ks_list, vs_list, s_out_list = self.extract_feature(
+                    style_latent,
+                    t,
+                    null_embeds_for_style,
+                )
+                if content_image is not None:
+                    qc_list, kc_list, vc_list, c_out_list = self.extract_feature(
+                        content_latent,
+                        t,
+                        null_embeds_for_content,
+                    )
+            for j in range(iters):
+                style_loss = 0
+                content_loss = 0
+                optimizer.zero_grad()
+                q_list, k_list, v_list, self_out_list = self.extract_feature(
+                    latents,
+                    t,
+                    null_embeds_for_latents,
+                )
+                style_loss = ad_loss(q_list, ks_list, vs_list, self_out_list, scale=attn_scale, source_mask=source_mask, target_mask=target_mask)
+                if content_image is not None:
+                    content_loss = q_loss(q_list, qc_list)
+                    # content_loss = qk_loss(q_list, k_list, qc_list, kc_list)
+                    # content_loss = qkv_loss(q_list, k_list, vc_list, c_out_list)
+                loss = style_loss + content_loss * weight
+                self.accelerator.backward(loss)
+                optimizer.step()
+                pbar.set_postfix(loss=loss.item(), time=t.item(), iter=j)
+        images = self.latent2image(latents)
+        # Offload all models
+        self.maybe_free_model_hooks()
+        return images
+
+    def panorama(
+        self,
+        lr=0.05,
+        iters=1,
+        attn_scale=1,
+        adain=False,
+        controller=None,
+        style_image=None,
+        mixed_precision="no",
+        enable_gradient_checkpoint=False,
+        prompt: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 1,
+        stride=8,
+        view_batch_size: int = 16,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        guidance_rescale: float = 0.0,
+        clip_skip: Optional[int] = None,
+        **kwargs,
+    ):
+
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        self._guidance_scale = guidance_scale
+        self._guidance_rescale = guidance_rescale
+        self._clip_skip = clip_skip
+        self._cross_attention_kwargs = cross_attention_kwargs
+        self._interrupt = False
+
+        self.accelerator = Accelerator(
+            mixed_precision=mixed_precision, gradient_accumulation_steps=1
+        )
+        self.init(enable_gradient_checkpoint)
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_cfg = guidance_scale > 1.0
+
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+                self.do_classifier_free_guidance,
+            )
+
+        # 3. Encode input prompt
+        text_encoder_lora_scale = (
+            cross_attention_kwargs.get("scale", None)
+            if cross_attention_kwargs is not None
+            else None
+        )
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            do_cfg,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            lora_scale=text_encoder_lora_scale,
+            clip_skip=clip_skip,
+        )
+        # For classifier free guidance, we need to do two forward passes.
+        # Here we concatenate the unconditional and text embeddings into a single batch
+        # to avoid doing two forward passes
+        if do_cfg:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Define panorama grid and initialize views for synthesis.
+        # prepare batch grid
+        views = self.get_views_(height, width, window_size=64, stride=stride)
+        views_batch = [
+            views[i : i + view_batch_size]
+            for i in range(0, len(views), view_batch_size)
+        ]
+        print(len(views), len(views_batch), views_batch)
+        self.scheduler.set_timesteps(num_inference_steps)
+        views_scheduler_status = [copy.deepcopy(self.scheduler.__dict__)] * len(
+            views_batch
+        )
+        count = torch.zeros_like(latents)
+        value = torch.zeros_like(latents)
+
+        # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7.1 Add image embeds for IP-Adapter
+        added_cond_kwargs = (
+            {"image_embeds": image_embeds}
+            if ip_adapter_image is not None or ip_adapter_image_embeds is not None
+            else None
+        )
+
+        # 7.2 Optionally get Guidance Scale Embedding
+        timestep_cond = None
+        if self.unet.config.time_cond_proj_dim is not None:
+            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(
+                batch_size * num_images_per_prompt
+            )
+            timestep_cond = self.get_guidance_scale_embedding(
+                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
+            ).to(device=device, dtype=latents.dtype)
+
+        # 8. Denoising loop
+        # Each denoising step also includes refinement of the latents with respect to the
+        # views.
+
+        timesteps = self.scheduler.timesteps
+        self.style_latent = self.image2latent(style_image)
+        self.content_latent = None
+        null_embeds = self.encode_prompt("", device, 1, False)[0]
+        self.null_embeds = null_embeds
+        self.null_embeds_for_latents = torch.cat([null_embeds] * latents.shape[0])
+        self.null_embeds_for_style = torch.cat(
+            [null_embeds] * self.style_latent.shape[0]
+        )
+        self.adain = adain
+        self.attn_scale = attn_scale
+        self.cache = utils.DataCache()
+        self.controller = controller
+        utils.register_attn_control(
+            self.classifier, controller=self.controller, cache=self.cache
+        )
+        print("Total self attention layers of Unet: ", controller.num_self_layers)
+        print("Self attention layers for AD: ", controller.self_layers)
+
+        pbar = tqdm(timesteps, desc="Sample")
+        for i, t in enumerate(pbar):
+            count.zero_()
+            value.zero_()
+            # generate views
+            # Here, we iterate through different spatial crops of the latents and denoise them. These
+            # denoised (latent) crops are then averaged to produce the final latent
+            # for the current timestep via MultiDiffusion. Please see Sec. 4.1 in the
+            # MultiDiffusion paper for more details: https://arxiv.org/abs/2302.08113
+            # Batch views denoise
+            for j, batch_view in enumerate(views_batch):
+                vb_size = len(batch_view)
+                # get the latents corresponding to the current view coordinates
+                latents_for_view = torch.cat(
+                    [
+                        latents[:, :, h_start:h_end, w_start:w_end]
+                        for h_start, h_end, w_start, w_end in batch_view
+                    ]
+                )
+                # rematch block's scheduler status
+                self.scheduler.__dict__.update(views_scheduler_status[j])
+
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = (
+                    latents_for_view.repeat_interleave(2, dim=0)
+                    if do_cfg
+                    else latents_for_view
+                )
+
+                latent_model_input = self.scheduler.scale_model_input(
+                    latent_model_input, t
+                )
+
+                # repeat prompt_embeds for batch
+                prompt_embeds_input = torch.cat([prompt_embeds] * vb_size)
+
+                # predict the noise residual
+                with torch.no_grad():
+                    noise_pred = self.unet(
+                        latent_model_input,
+                        t,
+                        encoder_hidden_states=prompt_embeds_input,
+                        timestep_cond=timestep_cond,
+                        cross_attention_kwargs=cross_attention_kwargs,
+                        added_cond_kwargs=added_cond_kwargs,
+                    ).sample
+
+                    # perform guidance
+                    if do_cfg:
+                        noise_pred_uncond, noise_pred_text = (
+                            noise_pred[::2],
+                            noise_pred[1::2],
+                        )
+                        noise_pred = noise_pred_uncond + guidance_scale * (
+                            noise_pred_text - noise_pred_uncond
+                        )
+
+                    # compute the previous noisy sample x_t -> x_t-1
+                    latents_denoised_batch = self.scheduler.step(
+                        noise_pred, t, latents_for_view, **extra_step_kwargs
+                    ).prev_sample
+                if iters > 0:
+                    self.null_embeds_for_latents = torch.cat(
+                        [self.null_embeds] * noise_pred.shape[0]
+                    )
+                    latents_denoised_batch = self.AD(
+                        latents_denoised_batch, t, lr, iters, pbar
+                    )
+                # save views scheduler status after sample
+                views_scheduler_status[j] = copy.deepcopy(self.scheduler.__dict__)
+
+                # extract value from batch
+                for latents_view_denoised, (h_start, h_end, w_start, w_end) in zip(
+                    latents_denoised_batch.chunk(vb_size), batch_view
+                ):
+
+                    value[:, :, h_start:h_end, w_start:w_end] += latents_view_denoised
+                    count[:, :, h_start:h_end, w_start:w_end] += 1
+
+            # take the MultiDiffusion step. Eq. 5 in MultiDiffusion paper: https://arxiv.org/abs/2302.08113
+            latents = torch.where(count > 0, value / count, value)
+
+        images = self.latent2image(latents)
+        # Offload all models
+        self.maybe_free_model_hooks()
+        return images
+
+    def AD(self, latents, t, lr, iters, pbar, weight=0):
+        t = max(
+            t
+            - self.scheduler.config.num_train_timesteps
+            // self.scheduler.num_inference_steps,
+            torch.tensor([0], device=self.device),
+        )
+        if self.adain:
+            noise = torch.randn_like(self.style_latent)
+            style_latent = self.scheduler.add_noise(self.style_latent, noise, t)
+            latents = utils.adain(latents, style_latent)
+
+        with torch.no_grad():
+            qs_list, ks_list, vs_list, s_out_list = self.extract_feature(
+                self.style_latent,
+                t,
+                self.null_embeds_for_style,
+                add_noise=True,
+            )
+            if self.content_latent is not None:
+                qc_list, kc_list, vc_list, c_out_list = self.extract_feature(
+                    self.content_latent,
+                    t,
+                    self.null_embeds,
+                    add_noise=True,
+                )
+
+        latents = latents.detach()
+        optimizer = torch.optim.Adam([latents.requires_grad_()], lr=lr)
+        optimizer = self.accelerator.prepare(optimizer)
+
+        for j in range(iters):
+            style_loss = 0
+            content_loss = 0
+            optimizer.zero_grad()
+            q_list, k_list, v_list, self_out_list = self.extract_feature(
+                latents,
+                t,
+                self.null_embeds_for_latents,
+                add_noise=False,
+            )
+            style_loss = ad_loss(q_list, ks_list, vs_list, self_out_list, scale=self.attn_scale)
+            if self.content_latent is not None:
+                content_loss = q_loss(q_list, qc_list)
+                # content_loss = qk_loss(q_list, k_list, qc_list, kc_list)
+                # content_loss = qkv_loss(q_list, k_list, vc_list, c_out_list)
+            loss = style_loss + content_loss * weight
+            self.accelerator.backward(loss)
+            optimizer.step()
+
+            pbar.set_postfix(loss=loss.item(), time=t.item(), iter=j)
+        latents = latents.detach()
+        return latents
+
+    def extract_feature(
+        self,
+        latent,
+        t,
+        embeds,
+        add_noise=False,
+    ):
+        self.cache.clear()
+        self.controller.step()
+        if add_noise:
+            noise = torch.randn_like(latent)
+            latent_ = self.scheduler.add_noise(latent, noise, t)
+        else:
+            latent_ = latent
+        _ = self.classifier(latent_, t, embeds)[0]
+        return self.cache.get()
+
+    def get_views_(
+        self,
+        panorama_height: int,
+        panorama_width: int,
+        window_size: int = 64,
+        stride: int = 8,
+    ) -> List[Tuple[int, int, int, int]]:
+        panorama_height //= 8
+        panorama_width //= 8
+
+        num_blocks_height = (
+            math.ceil((panorama_height - window_size) / stride) + 1
+            if panorama_height > window_size
+            else 1
+        )
+        num_blocks_width = (
+            math.ceil((panorama_width - window_size) / stride) + 1
+            if panorama_width > window_size
+            else 1
+        )
+
+        views = []
+        for i in range(int(num_blocks_height)):
+            for j in range(int(num_blocks_width)):
+                h_start = int(min(i * stride, panorama_height - window_size))
+                w_start = int(min(j * stride, panorama_width - window_size))
+
+                h_end = h_start + window_size
+                w_end = w_start + window_size
+
+                views.append((h_start, h_end, w_start, w_end))
+
+        return views

pipeline_sdxl.py

@@ -0,0 +1,403 @@
+import math
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import utils
+from accelerate import Accelerator
+from accelerate.utils import (
+    DistributedDataParallelKwargs,
+    ProjectConfiguration,
+    set_seed,
+)
+from diffusers import StableDiffusionXLPipeline
+from diffusers.image_processor import PipelineImageInput
+from diffusers.utils.torch_utils import is_compiled_module
+from losses import *
+
+# from peft import LoraConfig, set_peft_model_state_dict
+from tqdm import tqdm
+
+
+class ADPipeline(StableDiffusionXLPipeline):
+    def freeze(self):
+        self.unet.requires_grad_(False)
+        self.text_encoder.requires_grad_(False)
+        self.text_encoder_2.requires_grad_(False)
+        self.vae.requires_grad_(False)
+        self.classifier.requires_grad_(False)
+
+    @torch.no_grad()
+    def image2latent(self, image):
+        dtype = next(self.vae.parameters()).dtype
+        device = self._execution_device
+        image = image.to(device=device, dtype=dtype) * 2.0 - 1.0
+        latent = self.vae.encode(image)["latent_dist"].mean
+        latent = latent * self.vae.config.scaling_factor
+        return latent
+
+    @torch.no_grad()
+    def latent2image(self, latent):
+        dtype = next(self.vae.parameters()).dtype
+        device = self._execution_device
+        latent = latent.to(device=device, dtype=dtype)
+        latent = latent / self.vae.config.scaling_factor
+        image = self.vae.decode(latent)[0]
+        return (image * 0.5 + 0.5).clamp(0, 1)
+
+    def init(self, enable_gradient_checkpoint):
+        self.freeze()
+        weight_dtype = torch.float32
+        if self.accelerator.mixed_precision == "fp16":
+            weight_dtype = torch.float16
+        elif self.accelerator.mixed_precision == "bf16":
+            weight_dtype = torch.bfloat16
+
+        # Move unet, vae and text_encoder to device and cast to weight_dtype
+        self.unet.to(self.accelerator.device, dtype=weight_dtype)
+        self.vae.to(self.accelerator.device, dtype=weight_dtype)
+        self.text_encoder.to(self.accelerator.device, dtype=weight_dtype)
+        self.text_encoder_2.to(self.accelerator.device, dtype=weight_dtype)
+        self.classifier.to(self.accelerator.device, dtype=weight_dtype)
+        self.classifier = self.accelerator.prepare(self.classifier)
+        if enable_gradient_checkpoint:
+            self.classifier.enable_gradient_checkpointing()
+            # self.classifier.train()
+     
+
+    def sample(
+        self,
+        lr=0.05,
+        iters=1,
+        adain=True,
+        controller=None,
+        style_image=None,
+        mixed_precision="no",
+        init_from_style=False,
+        start_time=999,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        denoising_end: Optional[float] = None,
+        guidance_scale: float = 5.0,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        pooled_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        guidance_rescale: float = 0.0,
+        original_size: Optional[Tuple[int, int]] = None,
+        crops_coords_top_left: Tuple[int, int] = (0, 0),
+        target_size: Optional[Tuple[int, int]] = None,
+        negative_original_size: Optional[Tuple[int, int]] = None,
+        negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
+        negative_target_size: Optional[Tuple[int, int]] = None,
+        clip_skip: Optional[int] = None,
+        enable_gradient_checkpoint=False,
+        **kwargs,
+    ):
+        # 0. Default height and width to unet
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        original_size = original_size or (height, width)
+        target_size = target_size or (height, width)
+        self._guidance_scale = guidance_scale
+        self._guidance_rescale = guidance_rescale
+        self._clip_skip = clip_skip
+        self._cross_attention_kwargs = cross_attention_kwargs
+        self._denoising_end = denoising_end
+        self._interrupt = False
+
+        self.accelerator = Accelerator(
+            mixed_precision=mixed_precision, gradient_accumulation_steps=1
+        )
+        self.init(enable_gradient_checkpoint)
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        # 3. Encode input prompt
+        lora_scale = (
+            self.cross_attention_kwargs.get("scale", None)
+            if self.cross_attention_kwargs is not None
+            else None
+        )
+
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            lora_scale=lora_scale,
+            clip_skip=self.clip_skip,
+        )
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 7. Prepare added time ids & embeddings
+        add_text_embeds = pooled_prompt_embeds
+        if self.text_encoder_2 is None:
+            text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
+        else:
+            text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
+
+        add_time_ids = self._get_add_time_ids(
+            original_size,
+            crops_coords_top_left,
+            target_size,
+            dtype=prompt_embeds.dtype,
+            text_encoder_projection_dim=text_encoder_projection_dim,
+        )
+        null_add_time_ids = add_time_ids.to(device)
+        if negative_original_size is not None and negative_target_size is not None:
+            negative_add_time_ids = self._get_add_time_ids(
+                negative_original_size,
+                negative_crops_coords_top_left,
+                negative_target_size,
+                dtype=prompt_embeds.dtype,
+                text_encoder_projection_dim=text_encoder_projection_dim,
+            )
+        else:
+            negative_add_time_ids = add_time_ids
+
+        if self.do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            add_text_embeds = torch.cat(
+                [negative_pooled_prompt_embeds, add_text_embeds], dim=0
+            )
+            add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
+
+        prompt_embeds = prompt_embeds.to(device)
+        add_text_embeds = add_text_embeds.to(device)
+        add_time_ids = add_time_ids.to(device).repeat(
+            batch_size * num_images_per_prompt, 1
+        )
+
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+                self.do_classifier_free_guidance,
+            )
+        # 8.1 Apply denoising_end
+        if (
+            self.denoising_end is not None
+            and isinstance(self.denoising_end, float)
+            and self.denoising_end > 0
+            and self.denoising_end < 1
+        ):
+            discrete_timestep_cutoff = int(
+                round(
+                    self.scheduler.config.num_train_timesteps
+                    - (self.denoising_end * self.scheduler.config.num_train_timesteps)
+                )
+            )
+            num_inference_steps = len(
+                list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))
+            )
+            timesteps = timesteps[:num_inference_steps]
+
+        # 9. Optionally get Guidance Scale Embedding
+        timestep_cond = None
+        if self.unet.config.time_cond_proj_dim is not None:
+            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(
+                batch_size * num_images_per_prompt
+            )
+            timestep_cond = self.get_guidance_scale_embedding(
+                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
+            ).to(device=device, dtype=latents.dtype)
+        self.timestep_cond = timestep_cond
+        (null_embeds, _, null_pooled_embeds, _) = self.encode_prompt("", device=device)
+
+        added_cond_kwargs = {
+            "text_embeds": add_text_embeds, 
+            "time_ids": add_time_ids
+            }
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            added_cond_kwargs["image_embeds"] = image_embeds
+
+        self.scheduler.set_timesteps(num_inference_steps)
+
+        timesteps = self.scheduler.timesteps
+        style_latent = self.image2latent(style_image)
+        if init_from_style:
+            latents = torch.cat([style_latent] * latents.shape[0])
+            noise = torch.randn_like(latents)
+            latents = self.scheduler.add_noise(
+                latents,
+                noise,
+                torch.tensor([999]),
+            )
+
+        self.style_latent = style_latent
+        self.null_embeds_for_latents = torch.cat([null_embeds] * (latents.shape[0]))
+        self.null_embeds_for_style = torch.cat([null_embeds] * style_latent.shape[0])
+        self.null_added_cond_kwargs_for_latents = {
+            "text_embeds": torch.cat([null_pooled_embeds] * (latents.shape[0])),
+            "time_ids": torch.cat([null_add_time_ids] * (latents.shape[0])),
+        }
+        self.null_added_cond_kwargs_for_style = {
+            "text_embeds": torch.cat([null_pooled_embeds] * style_latent.shape[0]),
+            "time_ids": torch.cat([null_add_time_ids] * style_latent.shape[0]),
+        }
+        self.adain = adain
+        self.cache = utils.DataCache()
+        self.controller = controller
+        utils.register_attn_control(
+            self.classifier, controller=controller, cache=self.cache
+        )
+        print("Total self attention layers of Unet: ", controller.num_self_layers)
+        print("Self attention layers for AD: ", controller.self_layers)
+
+        pbar = tqdm(timesteps, desc="Sample")
+        for i, t in enumerate(pbar):
+            with torch.no_grad():
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = (
+                    torch.cat([latents] * 2)
+                    if self.do_classifier_free_guidance
+                    else latents
+                )
+
+                # predict the noise residual
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    timestep_cond=timestep_cond,
+                    cross_attention_kwargs=self.cross_attention_kwargs,
+                    added_cond_kwargs=added_cond_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # perform guidance
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (
+                        noise_pred_text - noise_pred_uncond
+                    )
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                
+            if iters > 0 and t < start_time:
+                latents = self.AD(latents, t, lr, iters, pbar)    
+                
+                
+        # Offload all models
+        # self.enable_model_cpu_offload()
+        images = self.latent2image(latents)
+        self.maybe_free_model_hooks()
+        return images
+
+    def AD(self, latents, t, lr, iters, pbar):
+        t = max(
+            t
+            - self.scheduler.config.num_train_timesteps
+            // self.scheduler.num_inference_steps,
+            torch.tensor([0], device=self.device),
+        )
+
+        if self.adain:
+            noise = torch.randn_like(self.style_latent)
+            style_latent = self.scheduler.add_noise(self.style_latent, noise, t)
+            latents = utils.adain(latents, style_latent)
+
+        with torch.no_grad():
+            qs_list, ks_list, vs_list, s_out_list = self.extract_feature(
+                self.style_latent,
+                t,
+                self.null_embeds_for_style,
+                self.timestep_cond,
+                self.null_added_cond_kwargs_for_style,
+                add_noise=True,
+            )
+        # latents = latents.to(dtype=torch.float32)
+        latents = latents.detach()
+        optimizer = torch.optim.Adam([latents.requires_grad_()], lr=lr)
+        optimizer, latents = self.accelerator.prepare(optimizer, latents)
+
+        for j in range(iters):
+            optimizer.zero_grad()
+            q_list, k_list, v_list, self_out_list = self.extract_feature(
+                latents,
+                t,
+                self.null_embeds_for_latents,
+                self.timestep_cond,
+                self.null_added_cond_kwargs_for_latents,
+                add_noise=False,
+            )
+
+            loss = ad_loss(q_list, ks_list, vs_list, self_out_list)
+            self.accelerator.backward(loss)
+            optimizer.step()
+
+            pbar.set_postfix(loss=loss.item(), time=t.item(), iter=j)
+        latents = latents.detach()
+        return latents
+
+    def extract_feature(
+        self,
+        latent,
+        t,
+        encoder_hidden_states,
+        timestep_cond,
+        added_cond_kwargs,
+        add_noise=False,
+    ):
+        self.cache.clear()
+        self.controller.step()
+        if add_noise:
+            noise = torch.randn_like(latent)
+            latent_ = self.scheduler.add_noise(latent, noise, t)
+        else:
+            latent_ = latent
+        self.classifier(
+            latent_,
+            t,
+            encoder_hidden_states=encoder_hidden_states,
+            timestep_cond=timestep_cond,
+            added_cond_kwargs=added_cond_kwargs,
+            return_dict=False,
+        )[0]
+        return self.cache.get()

requirements.txt

@@ -0,0 +1,10 @@
+diffusers
+torch>=2.0.0
+torchvision
+transformers
+accelerate
+safetensors
+spaces
+huggingface-hub
+gradio
+matplotlib

train_vae.py

@@ -0,0 +1,87 @@
+import argparse
+import os
+
+import torch
+from diffusers import AutoencoderKL
+from torch import nn
+from torch.optim import Adam
+from utils import load_image, save_image
+
+
+def main(args):
+    os.makedirs(args.out_dir, exist_ok=True)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    vae = AutoencoderKL.from_pretrained(args.vae_model_path).to(
+        device, dtype=torch.float32
+    )
+    vae.requires_grad_(False)
+
+    image = load_image(args.image_path, size=(512, 512)).to(device, dtype=torch.float32)
+    image = image * 2 - 1
+    save_image(image / 2 + 0.5, f"{args.out_dir}/ori_image.png")
+
+    latents = vae.encode(image)["latent_dist"].mean
+    save_image(latents, f"{args.out_dir}/latents.png")
+
+    rec_image = vae.decode(latents, return_dict=False)[0]
+    save_image(rec_image / 2 + 0.5, f"{args.out_dir}/rec_image.png")
+
+    for param in vae.decoder.parameters():
+        param.requires_grad = True
+
+    loss_fn = nn.L1Loss()
+    optimizer = Adam(vae.decoder.parameters(), lr=args.learning_rate)
+
+    # Training loop
+    for epoch in range(args.num_epochs):
+        reconstructed = vae.decode(latents, return_dict=False)[0]
+        loss = loss_fn(reconstructed, image)
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        print(f"Epoch {epoch+1}/{args.num_epochs}, Loss: {loss.item()}")
+
+    rec_image = vae.decode(latents, return_dict=False)[0]
+    save_image(rec_image / 2 + 0.5, f"{args.out_dir}/trained_rec_image.png")
+    vae.save_pretrained(
+        f"{args.out_dir}/trained_vae_{os.path.basename(args.image_path)}"
+    )
+
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Train a VAE with given image and settings."
+    )
+
+    # Add arguments
+    parser.add_argument(
+        "--out_dir",
+        type=str,
+        default="./trained_vae/",
+        help="Output directory to save results",
+    )
+    parser.add_argument(
+        "--vae_model_path",
+        type=str,
+        required=True,
+        help="Path to the pretrained VAE model",
+    )
+    parser.add_argument(
+        "--image_path", type=str, required=True, help="Path to the input image"
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=1e-4,
+        help="Learning rate for the optimizer",
+    )
+    parser.add_argument(
+        "--num_epochs", type=int, default=75, help="Number of training epochs"
+    )
+
+    args = parser.parse_args()
+    main(args)

utils.py

@@ -0,0 +1,195 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from torchvision.transforms import ToTensor
+from torchvision.utils import save_image
+import matplotlib.pyplot as plt
+import math
+
+
+def register_attn_control(unet, controller, cache=None):
+    def attn_forward(self):
+        def forward(
+            hidden_states,
+            encoder_hidden_states=None,
+            attention_mask=None,
+            temb=None,
+            *args,
+            **kwargs,
+        ):
+            residual = hidden_states
+            if self.spatial_norm is not None:
+                hidden_states = self.spatial_norm(hidden_states, temb)
+
+            input_ndim = hidden_states.ndim
+
+            if input_ndim == 4:
+                batch_size, channel, height, width = hidden_states.shape
+                hidden_states = hidden_states.view(
+                    batch_size, channel, height * width
+                ).transpose(1, 2)
+
+            batch_size, sequence_length, _ = (
+                hidden_states.shape
+                if encoder_hidden_states is None
+                else encoder_hidden_states.shape
+            )
+
+            if attention_mask is not None:
+                attention_mask = self.prepare_attention_mask(
+                    attention_mask, sequence_length, batch_size
+                )
+                # scaled_dot_product_attention expects attention_mask shape to be
+                # (batch, heads, source_length, target_length)
+                attention_mask = attention_mask.view(
+                    batch_size, self.heads, -1, attention_mask.shape[-1]
+                )
+
+            if self.group_norm is not None:
+                hidden_states = self.group_norm(
+                    hidden_states.transpose(1, 2)
+                ).transpose(1, 2)
+
+            q = self.to_q(hidden_states)
+            is_self = encoder_hidden_states is None
+
+            if encoder_hidden_states is None:
+                encoder_hidden_states = hidden_states
+            elif self.norm_cross:
+                encoder_hidden_states = self.norm_encoder_hidden_states(
+                    encoder_hidden_states
+                )
+
+            k = self.to_k(encoder_hidden_states)
+            v = self.to_v(encoder_hidden_states)
+
+            inner_dim = k.shape[-1]
+            head_dim = inner_dim // self.heads
+
+            q = q.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+            k = k.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+            v = v.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+            # the output of sdp = (batch, num_heads, seq_len, head_dim)
+            # TODO: add support for attn.scale when we move to Torch 2.1
+            hidden_states = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+            )
+            if is_self and controller.cur_self_layer in controller.self_layers:
+                cache.add(q, k, v, hidden_states)
+
+            hidden_states = hidden_states.transpose(1, 2).reshape(
+                batch_size, -1, self.heads * head_dim
+            )
+            hidden_states = hidden_states.to(q.dtype)
+
+            # linear proj
+            hidden_states = self.to_out[0](hidden_states)
+            # dropout
+            hidden_states = self.to_out[1](hidden_states)
+
+            if input_ndim == 4:
+                hidden_states = hidden_states.transpose(-1, -2).reshape(
+                    batch_size, channel, height, width
+                )
+            if self.residual_connection:
+                hidden_states = hidden_states + residual
+
+            hidden_states = hidden_states / self.rescale_output_factor
+
+            if is_self:
+                controller.cur_self_layer += 1
+
+            return hidden_states
+
+        return forward
+
+    def modify_forward(net, count):
+        for name, subnet in net.named_children():
+            if net.__class__.__name__ == "Attention":  # spatial Transformer layer
+                net.forward = attn_forward(net)
+                return count + 1
+            elif hasattr(net, "children"):
+                count = modify_forward(subnet, count)
+        return count
+
+    cross_att_count = 0
+    for net_name, net in unet.named_children():
+        cross_att_count += modify_forward(net, 0)
+    controller.num_self_layers = cross_att_count // 2
+
+
+def load_image(image_path, size=None, mode="RGB"):
+    img = Image.open(image_path).convert(mode)
+    if size is None:
+        width, height = img.size
+        new_width = (width // 64) * 64
+        new_height = (height // 64) * 64
+        size = (new_width, new_height)
+    img = img.resize(size, Image.BICUBIC)
+    return ToTensor()(img).unsqueeze(0)
+
+
+def adain(source, target, eps=1e-6):
+    source_mean, source_std = torch.mean(source, dim=(2, 3), keepdim=True), torch.std(
+        source, dim=(2, 3), keepdim=True
+    )
+    target_mean, target_std = torch.mean(
+        target, dim=(0, 2, 3), keepdim=True
+    ), torch.std(target, dim=(0, 2, 3), keepdim=True)
+    normalized_source = (source - source_mean) / (source_std + eps)
+    transferred_source = normalized_source * target_std + target_mean
+
+    return transferred_source
+
+
+class Controller:
+    def step(self):
+        self.cur_self_layer = 0
+
+    def __init__(self, self_layers=(0, 16)):
+        self.num_self_layers = -1
+        self.cur_self_layer = 0
+        self.self_layers = list(range(*self_layers))
+
+
+class DataCache:
+    def __init__(self):
+        self.q = []
+        self.k = []
+        self.v = []
+        self.out = []
+
+    def clear(self):
+        self.q.clear()
+        self.k.clear()
+        self.v.clear()
+        self.out.clear()
+
+    def add(self, q, k, v, out):
+        self.q.append(q)
+        self.k.append(k)
+        self.v.append(v)
+        self.out.append(out)
+
+    def get(self):
+        return self.q.copy(), self.k.copy(), self.v.copy(), self.out.copy()
+    
+
+
+def show_image(path, title, display_height=3, title_fontsize=12):
+    img = Image.open(path)
+    img_width, img_height = img.size
+
+    aspect_ratio = img_width / img_height
+    display_width = display_height * aspect_ratio
+
+    plt.figure(figsize=(display_width, display_height))
+    plt.imshow(img)
+    plt.title(title, 
+             fontsize=title_fontsize, 
+             fontweight='bold', 
+             pad=20) 
+    plt.axis('off')    
+    plt.tight_layout() 
+    plt.show()

webui/__init__.py

@@ -0,0 +1,5 @@
+
+from .tab_style_t2i import create_interface_style_t2i
+from .tab_style_transfer import create_interface_style_transfer
+from .tab_texture_synthesis import create_interface_texture_synthesis
+from .runner import Runner

webui/runner.py

@@ -0,0 +1,157 @@
+import torch
+from PIL import Image
+from diffusers import DDIMScheduler
+from accelerate.utils import set_seed
+from torchvision.transforms.functional import to_pil_image, to_tensor
+
+from pipeline_sd import ADPipeline
+from pipeline_sdxl import ADPipeline as ADXLPipeline
+from utils import Controller
+
+import os
+
+
+class Runner:
+    def __init__(self):
+        self.sd15 = None
+        self.sdxl = None
+        self.loss_fn = torch.nn.L1Loss(reduction="mean")
+    
+    def load_pipeline(self, model_path_or_name):
+
+        if 'xl' in model_path_or_name and self.sdxl is None:
+            scheduler = DDIMScheduler.from_pretrained(os.path.join('./checkpoints', model_path_or_name), subfolder="scheduler")
+            self.sdxl = ADXLPipeline.from_pretrained(os.path.join('./checkpoints', model_path_or_name), scheduler=scheduler, safety_checker=None)
+            self.sdxl.classifier = self.sdxl.unet
+        elif self.sd15 is None:
+            scheduler = DDIMScheduler.from_pretrained(os.path.join('./checkpoints', model_path_or_name), subfolder="scheduler")
+            self.sd15 = ADPipeline.from_pretrained(os.path.join('./checkpoints', model_path_or_name), scheduler=scheduler, safety_checker=None)
+            self.sd15.classifier = self.sd15.unet
+
+    def preprocecss(self, image: Image.Image, height=None, width=None):
+        if width is None or height is None: 
+            width, height = image.size
+        new_width = (width // 64) * 64
+        new_height = (height // 64) * 64
+        size = (new_width, new_height)
+        image = image.resize(size, Image.BICUBIC)
+        return to_tensor(image).unsqueeze(0)
+
+    def run_style_transfer(self, content_image, style_image, seed, num_steps, lr, content_weight, mixed_precision, model, **kwargs):
+        self.load_pipeline(model)
+
+        content_image = self.preprocecss(content_image)
+        style_image = self.preprocecss(style_image, height=512, width=512)
+
+        height, width = content_image.shape[-2:]
+        set_seed(seed)
+        controller = Controller(self_layers=(10, 16))
+        result = self.sd15.optimize(
+            lr=lr,
+            batch_size=1,
+            iters=1,
+            width=width,
+            height=height,
+            weight=content_weight,
+            controller=controller,
+            style_image=style_image,
+            content_image=content_image,
+            mixed_precision=mixed_precision,
+            num_inference_steps=num_steps,
+            enable_gradient_checkpoint=False,
+        )
+        output_image = to_pil_image(result[0])
+        del result
+        torch.cuda.empty_cache()
+        return [output_image]
+
+    def run_style_t2i_generation(self, style_image, prompt, negative_prompt, guidance_scale, height, width, seed, num_steps, iterations, lr, num_images_per_prompt, mixed_precision, is_adain, model):
+        self.load_pipeline(model)
+
+        use_xl = 'xl' in model
+        height, width = (1024, 1024) if 'xl' in model else (512, 512)
+        style_image = self.preprocecss(style_image, height=height, width=width)
+
+        set_seed(seed)
+        self_layers = (64, 70) if use_xl else (10, 16)
+        
+        controller = Controller(self_layers=self_layers)
+
+        pipeline = self.sdxl if use_xl else self.sd15
+        images = pipeline.sample(
+            controller=controller,
+            iters=iterations,
+            lr=lr,
+            adain=is_adain,
+            height=height,
+            width=width,
+            mixed_precision=mixed_precision,
+            style_image=style_image,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_steps,
+            num_images_per_prompt=num_images_per_prompt,
+            enable_gradient_checkpoint=False
+        )
+        output_images = [to_pil_image(image) for image in images]
+
+        del images
+        torch.cuda.empty_cache()
+        return output_images
+
+    def run_texture_synthesis(self, texture_image, height, width, seed, num_steps, iterations, lr, mixed_precision, num_images_per_prompt, synthesis_way,model):
+        self.load_pipeline(model) 
+
+        texture_image = self.preprocecss(texture_image, height=512, width=512)
+
+        set_seed(seed)
+        controller = Controller(self_layers=(10, 16))
+
+        if synthesis_way == 'Sampling':
+            results = self.sd15.sample(
+                lr=lr,
+                adain=False,
+                iters=iterations,
+                width=width,
+                height=height,
+                weight=0.,
+                controller=controller,
+                style_image=texture_image,
+                content_image=None,
+                prompt="",
+                negative_prompt="",
+                mixed_precision=mixed_precision,
+                num_inference_steps=num_steps,
+                guidance_scale=1.,
+                num_images_per_prompt=num_images_per_prompt,
+                enable_gradient_checkpoint=False,
+            )
+        elif synthesis_way == 'MultiDiffusion':   
+            results = self.sd15.panorama(
+                lr=lr,
+                iters=iterations,
+                width=width,
+                height=height,
+                weight=0.,
+                controller=controller,
+                style_image=texture_image,
+                content_image=None,
+                prompt="",
+                negative_prompt="",
+                stride=8,
+                view_batch_size=8,
+                mixed_precision=mixed_precision,
+                num_inference_steps=num_steps,
+                guidance_scale=1.,
+                num_images_per_prompt=num_images_per_prompt,
+                enable_gradient_checkpoint=False,
+            )
+        else:
+            raise ValueError
+        
+        output_images = [to_pil_image(image) for image in results]
+        del results
+        torch.cuda.empty_cache()
+        return output_images
+    

webui/tab_style_t2i.py

@@ -0,0 +1,51 @@
+import os
+from PIL import Image
+import gradio as gr
+
+
+def create_interface_style_t2i(runner):
+    with gr.Blocks():
+        with gr.Row():
+            gr.Markdown('1. Upload the style image and text your prompt.\n'
+                        '2. Choose the generative model.\n'
+                        '3. (Optional) Customize the configurations below as needed.\n'
+                        '4. Cilck `Run` to start generation.')
+        
+        with gr.Row():
+            with gr.Column():
+                style_image = gr.Image(label='Input Style Image', type='pil', interactive=True,
+                                        value=Image.open('examples/s1.jpg').convert('RGB') if os.path.exists('examples/s1.jpg') else None)
+                prompt = gr.Textbox(label='Prompt', value='A rocket')
+                negative_prompt = gr.Textbox(label='Negative Prompt', value='')
+
+                base_model_list = ['stable-diffusion-v1-5', 'stable-diffusion-xl-base-1.0']
+                model = gr.Radio(choices=base_model_list, label='Select a Base Model', value='stable-diffusion-xl-base-1.0')
+
+                run_button = gr.Button(value='Run')
+                
+                gr.Examples(
+                    [[Image.open('./webui/images/image_02_01.jpg').convert('RGB'), 'A rocket', 'stable-diffusion-xl-base-1.0']],
+                    [style_image, prompt, model]
+                )
+
+            with gr.Column():
+                with gr.Accordion('Options', open=True):
+                    guidance_scale = gr.Slider(label='Guidance Scale', minimum=1., maximum=30., value=7.5, step=0.1)
+                    height = gr.Number(label='Height', value=1024, precision=0, minimum=2, maximum=4096)
+                    width = gr.Number(label='Width', value=1024, precision=0, minimum=2, maximum=4096)
+                    seed = gr.Number(label='Seed', value=2025, precision=0, minimum=0, maximum=2**31)
+                    num_steps = gr.Slider(label='Number of Steps', minimum=1, maximum=1000, value=50, step=1)
+                    iterations = gr.Slider(label='Iterations', minimum=0, maximum=10, value=2, step=1)
+                    lr = gr.Slider(label='Learning Rate', minimum=0.01, maximum=0.5, value=0.015, step=0.001)
+                    num_images_per_prompt = gr.Slider(label='Num Images Per Prompt', minimum=1, maximum=10, value=1, step=1)
+                    mixed_precision = gr.Radio(choices=['bf16', 'no'], value='bf16', label='Mixed Precision')
+                    is_adain = gr.Checkbox(label='Adain', value=True,)
+                    
+            with gr.Column():
+                gr.Markdown('#### Output Image:\n')
+                result_gallery = gr.Gallery(label='Output', elem_id='gallery', columns=2, height='auto', preview=True)
+        
+        ips = [style_image, prompt, negative_prompt, guidance_scale, height, width, seed, num_steps, iterations, lr, num_images_per_prompt, mixed_precision, is_adain, model]
+
+        run_button.click(fn=runner.run_style_t2i_generation, inputs=ips, outputs=[result_gallery])
+

webui/tab_style_transfer.py

@@ -0,0 +1,45 @@
+import os
+from PIL import Image
+import gradio as gr
+
+
+def create_interface_style_transfer(runner):
+    with gr.Blocks():
+        with gr.Row():
+            gr.Markdown('1. Upload the content and style images as inputs.\n'
+                        '2. (Optional) Customize the configurations below as needed.\n'
+                        '3. Cilck `Run` to start transfer.')
+        
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    content_image = gr.Image(label='Input Content Image', type='pil', interactive=True,
+                                             value=Image.open('examples/c1.jpg').convert('RGB') if os.path.exists('examples/c1.jpg') else None)
+                    style_image = gr.Image(label='Input Style Image', type='pil', interactive=True,
+                                           value=Image.open('examples/s1.jpg').convert('RGB') if os.path.exists('examples/s1.jpg') else None)
+
+                run_button = gr.Button(value='Run')
+
+                with gr.Accordion('Options', open=True):
+                    seed = gr.Number(label='Seed', value=2025, precision=0, minimum=0, maximum=2**31)
+                    num_steps = gr.Slider(label='Number of Steps', minimum=1, maximum=1000, value=200, step=1)
+                    lr = gr.Slider(label='Learning Rate', minimum=0.01, maximum=0.5, value=0.05, step=0.01)
+                    content_weight = gr.Slider(label='Content Weight', minimum=0., maximum=1., value=0.25, step=0.001)
+                    mixed_precision = gr.Radio(choices=['bf16', 'no'], value='bf16', label='Mixed Precision')
+                    
+                    base_model_list = ['stable-diffusion-v1-5',]
+                    model = gr.Radio(choices=base_model_list, label='Select a Base Model', value='stable-diffusion-v1-5')
+
+            with gr.Column():
+                gr.Markdown('#### Output Image:\n')
+                result_gallery = gr.Gallery(label='Output', elem_id='gallery', columns=2, height='auto', preview=True)
+                
+                gr.Examples(
+                    [[Image.open('./webui/images/lecun.png').convert('RGB'), Image.open('./webui/images/40.jpg').convert('RGB'), 300, 0.23]],
+                    [content_image, style_image, num_steps, content_weight]
+                )
+
+        
+        ips = [content_image, style_image, seed, num_steps, lr, content_weight, mixed_precision, model]
+
+        run_button.click(fn=runner.run_style_transfer, inputs=ips, outputs=[result_gallery])

webui/tab_texture_synthesis.py

@@ -0,0 +1,46 @@
+import os
+from PIL import Image
+import gradio as gr
+
+
+def create_interface_texture_synthesis(runner):
+    with gr.Blocks():
+        with gr.Row():
+            gr.Markdown('1. Upload the texture image as input.\n'
+                        '2. (Optional) Customize the configurations below as needed.\n'
+                        '3. Cilck `Run` to start synthesis.')
+        
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    texture_image = gr.Image(label='Input Texture Image', type='pil', interactive=True,
+                                           value=Image.open('examples/s1.jpg').convert('RGB') if os.path.exists('examples/s1.jpg') else None)
+            
+                run_button = gr.Button(value='Run')
+
+                with gr.Accordion('Options', open=True):
+                    height = gr.Number(label='Height', value=512, precision=0, minimum=2, maximum=4096)
+                    width = gr.Number(label='Width', value=1024, precision=0, minimum=2, maximum=4096)
+                    seed = gr.Number(label='Seed', value=2025, precision=0, minimum=0, maximum=2**31)
+                    num_steps = gr.Slider(label='Number of Steps', minimum=1, maximum=1000, value=200, step=1)
+                    iterations = gr.Slider(label='Iterations', minimum=0, maximum=10, value=2, step=1)
+                    lr = gr.Slider(label='Learning Rate', minimum=0.01, maximum=0.5, value=0.05, step=0.01)
+                    mixed_precision = gr.Radio(choices=['bf16', 'no'], value='bf16', label='Mixed Precision')
+                    num_images_per_prompt = gr.Slider(label='Num Images Per Prompt', minimum=1, maximum=10, value=1, step=1)
+                    
+                    base_model_list = ['stable-diffusion-v1-5',]
+                    model = gr.Radio(choices=base_model_list, label='Select a Base Model', value='stable-diffusion-v1-5')
+                    synthesis_way = gr.Radio(['Sampling', 'MultiDiffusion'], label='Synthesis Way', value='MultiDiffusion')
+
+            with gr.Column():
+                gr.Markdown('#### Output Image:\n')
+                result_gallery = gr.Gallery(label='Output', elem_id='gallery', columns=2, height='auto', preview=True)
+
+                gr.Examples(
+                    [[Image.open('./webui/images/42.jpg').convert('RGB'), 'MultiDiffusion', 512, 1024]],
+                    [texture_image, synthesis_way, height, width]
+                )
+        ips = [texture_image, height, width, seed, num_steps, iterations, lr, mixed_precision, num_images_per_prompt, synthesis_way,model]
+
+        run_button.click(fn=runner.run_texture_synthesis, inputs=ips, outputs=[result_gallery])
+