Upload 13 files

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+doc/cn_example.jpg filter=lfs diff=lfs merge=lfs -text
+doc/md_example.jpg filter=lfs diff=lfs merge=lfs -text
+example_image/train.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,13 +1,37 @@
----
-title: Style-aligned Sdxl
-emoji: 🐨
-colorFrom: indigo
-colorTo: pink
-sdk: gradio
-sdk_version: 4.8.0
-app_file: app.py
-pinned: false
-license: mit
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Style Aligned Image Generation via Shared Attention
+
+
+### [Project Page](https://style-aligned-gen.github.io)   [Paper](https://style-aligned-gen.github.io/data/StyleAligned.pdf)
+
+
+## Setup
+
+This code was tested with Python 3.11, [Pytorch 2.1](https://pytorch.org/)  and [Diffusers 0.16](https://github.com/huggingface/diffusers).
+
+## Examples
+- See [**style_aligned_sdxl**][style_aligned] notebook for generating style aligned images using [SDXL](https://huggingface.co/docs/diffusers/using-diffusers/sdxl).
+
+![alt text](doc/sa_example.jpg)
+
+
+- See [**style_aligned_w_controlnet**][controlnet] notebook for generating style aligned and depth conditioned images using SDXL with [ControlNet-Depth](https://arxiv.org/abs/2302.05543).
+
+![alt text](doc/cn_example.jpg)
+
+
+-  [**style_aligned_w_multidiffusion**][multidiffusion] can be used for generating style aligned panoramas using [SD V2](https://huggingface.co/stabilityai/stable-diffusion-2) with [MultiDiffusion](https://multidiffusion.github.io/).
+
+![alt text](doc/md_example.jpg)
+
+## TODOs
+- [ ] Adding demo.
+- [ ] StyleAligned from an input image.
+- [ ] Multi-style with MultiDiffusion.
+- [ ] StyleAligned with DreamBooth
+
+## Disclaimer
+This is not an officially supported Google product.
+
+[style_aligned]: style_aligned_sdxl.ipynb
+[controlnet]: style_aligned_w_controlnet.ipynb
+[multidiffusion]: style_aligned_w_multidiffusion.ipynb

contributing.md

@@ -0,0 +1,28 @@
+# How to Contribute
+
+We'd love to accept your patches and contributions to this project. There are
+just a few small guidelines you need to follow.
+
+## Contributor License Agreement
+
+Contributions to this project must be accompanied by a Contributor License
+Agreement. You (or your employer) retain the copyright to your contribution;
+this simply gives us permission to use and redistribute your contributions as
+part of the project. Head over to <https://cla.developers.google.com/> to see
+your current agreements on file or to sign a new one.
+
+You generally only need to submit a CLA once, so if you've already submitted one
+(even if it was for a different project), you probably don't need to do it
+again.
+
+## Code Reviews
+
+All submissions, including submissions by project members, require review. We
+use GitHub pull requests for this purpose. Consult
+[GitHub Help](https://help.github.com/articles/about-pull-requests/) for more
+information on using pull requests.
+
+## Community Guidelines
+
+This project follows [Google's Open Source Community
+Guidelines](https://opensource.google/conduct/).

pipeline_calls.py

@@ -0,0 +1,552 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from __future__ import annotations
+from typing import Any
+import torch
+import numpy as np
+from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline
+from diffusers.image_processor import PipelineImageInput
+from diffusers.utils.torch_utils import is_compiled_module, is_torch_version
+from transformers import DPTImageProcessor, DPTForDepthEstimation
+from diffusers import StableDiffusionPanoramaPipeline
+from PIL import Image
+import copy
+
+T = torch.Tensor
+TN = T | None
+
+
+def get_depth_map(image: Image, feature_processor: DPTImageProcessor, depth_estimator: DPTForDepthEstimation) -> Image:
+    image = feature_processor(images=image, return_tensors="pt").pixel_values.to("cuda")
+    with torch.no_grad(), torch.autocast("cuda"):
+        depth_map = depth_estimator(image).predicted_depth
+
+    depth_map = torch.nn.functional.interpolate(
+        depth_map.unsqueeze(1),
+        size=(1024, 1024),
+        mode="bicubic",
+        align_corners=False,
+    )
+    depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
+    depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
+    depth_map = (depth_map - depth_min) / (depth_max - depth_min)
+    image = torch.cat([depth_map] * 3, dim=1)
+
+    image = image.permute(0, 2, 3, 1).cpu().numpy()[0]
+    image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8))
+    return image
+
+
+def concat_zero_control(control_reisduel: T) -> T:
+    b = control_reisduel.shape[0] // 2
+    zerso_reisduel = torch.zeros_like(control_reisduel[0:1])
+    return torch.cat((zerso_reisduel, control_reisduel[:b], zerso_reisduel, control_reisduel[b::]))
+
+
+@torch.no_grad()
+def controlnet_call(
+    pipeline: StableDiffusionXLControlNetPipeline,
+    prompt: str | list[str] = None,
+    prompt_2: str | list[str] | None = None,
+    image: PipelineImageInput = None,
+    height: int | None = None,
+    width: int | None = None,
+    num_inference_steps: int = 50,
+    guidance_scale: float = 5.0,
+    negative_prompt: str | list[str] | None = None,
+    negative_prompt_2: str | list[str] | None = None,
+    num_images_per_prompt: int = 1,
+    eta: float = 0.0,
+    generator: torch.Generator | None = None,
+    latents: TN = None,
+    prompt_embeds: TN = None,
+    negative_prompt_embeds: TN = None,
+    pooled_prompt_embeds: TN = None,
+    negative_pooled_prompt_embeds: TN = None,
+    cross_attention_kwargs: dict[str, Any] | None = None,
+    controlnet_conditioning_scale: float | list[float] = 1.0,
+    control_guidance_start: float | list[float] = 0.0,
+    control_guidance_end: float | list[float] = 1.0,
+    original_size: tuple[int, int] = None,
+    crops_coords_top_left: tuple[int, int] = (0, 0),
+    target_size: tuple[int, int] | None = None,
+    negative_original_size: tuple[int, int] | None = None,
+    negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+    negative_target_size:tuple[int, int] | None = None,
+    clip_skip: int | None = None,
+) -> list[Image]:
+    controlnet = pipeline.controlnet._orig_mod if is_compiled_module(pipeline.controlnet) else pipeline.controlnet
+
+    # align format for control guidance
+    if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list):
+        control_guidance_start = len(control_guidance_end) * [control_guidance_start]
+    elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list):
+        control_guidance_end = len(control_guidance_start) * [control_guidance_end]
+    elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list):
+        mult = 1
+        control_guidance_start, control_guidance_end = (
+            mult * [control_guidance_start],
+            mult * [control_guidance_end],
+        )
+
+    # 1. Check inputs. Raise error if not correct
+    pipeline.check_inputs(
+        prompt,
+        prompt_2,
+        image,
+        1,
+        negative_prompt,
+        negative_prompt_2,
+        prompt_embeds,
+        negative_prompt_embeds,
+        pooled_prompt_embeds,
+        negative_pooled_prompt_embeds,
+        controlnet_conditioning_scale,
+        control_guidance_start,
+        control_guidance_end,
+    )
+
+    pipeline._guidance_scale = guidance_scale
+
+    # 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 = pipeline._execution_device
+
+    # 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,
+        pooled_prompt_embeds,
+        negative_pooled_prompt_embeds,
+    ) = pipeline.encode_prompt(
+        prompt,
+        prompt_2,
+        device,
+        1,
+        True,
+        negative_prompt,
+        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=text_encoder_lora_scale,
+        clip_skip=clip_skip,
+    )
+
+    # 4. Prepare image
+    if isinstance(controlnet, ControlNetModel):
+        image = pipeline.prepare_image(
+            image=image,
+            width=width,
+            height=height,
+            batch_size=1,
+            num_images_per_prompt=1,
+            device=device,
+            dtype=controlnet.dtype,
+            do_classifier_free_guidance=True,
+            guess_mode=False,
+        )
+        height, width = image.shape[-2:]
+        image = torch.stack([image[0]] * num_images_per_prompt + [image[1]] * num_images_per_prompt)
+    else:
+        assert False
+    # 5. Prepare timesteps
+    pipeline.scheduler.set_timesteps(num_inference_steps, device=device)
+    timesteps = pipeline.scheduler.timesteps
+
+    # 6. Prepare latent variables
+    num_channels_latents = pipeline.unet.config.in_channels
+    latents = pipeline.prepare_latents(
+        1 + num_images_per_prompt,
+        num_channels_latents,
+        height,
+        width,
+        prompt_embeds.dtype,
+        device,
+        generator,
+        latents,
+    )
+        
+    # 6.5 Optionally get Guidance Scale Embedding
+    timestep_cond = None
+
+    # 7. Prepare extra step kwargs.
+    extra_step_kwargs = pipeline.prepare_extra_step_kwargs(generator, eta)
+
+    # 7.1 Create tensor stating which controlnets to keep
+    controlnet_keep = []
+    for i in range(len(timesteps)):
+        keeps = [
+            1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)
+            for s, e in zip(control_guidance_start, control_guidance_end)
+        ]
+        controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps)
+
+    # 7.2 Prepare added time ids & embeddings
+    if isinstance(image, list):
+        original_size = original_size or image[0].shape[-2:]
+    else:
+        original_size = original_size or image.shape[-2:]
+    target_size = target_size or (height, width)
+
+    add_text_embeds = pooled_prompt_embeds
+    if pipeline.text_encoder_2 is None:
+        text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
+    else:
+        text_encoder_projection_dim = pipeline.text_encoder_2.config.projection_dim
+
+    add_time_ids = pipeline._get_add_time_ids(
+        original_size,
+        crops_coords_top_left,
+        target_size,
+        dtype=prompt_embeds.dtype,
+        text_encoder_projection_dim=text_encoder_projection_dim,
+    )
+
+    if negative_original_size is not None and negative_target_size is not None:
+        negative_add_time_ids = pipeline._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
+
+    prompt_embeds = torch.stack([prompt_embeds[0]] + [prompt_embeds[1]] * num_images_per_prompt)
+    negative_prompt_embeds = torch.stack([negative_prompt_embeds[0]] + [negative_prompt_embeds[1]] * num_images_per_prompt)
+    negative_pooled_prompt_embeds = torch.stack([negative_pooled_prompt_embeds[0]] + [negative_pooled_prompt_embeds[1]] * num_images_per_prompt)
+    add_text_embeds = torch.stack([add_text_embeds[0]] + [add_text_embeds[1]] * num_images_per_prompt)
+    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(1 + num_images_per_prompt, 1)
+    batch_size = num_images_per_prompt + 1
+    # 8. Denoising loop
+    num_warmup_steps = len(timesteps) - num_inference_steps * pipeline.scheduler.order
+    is_unet_compiled = is_compiled_module(pipeline.unet)
+    is_controlnet_compiled = is_compiled_module(pipeline.controlnet)
+    is_torch_higher_equal_2_1 = is_torch_version(">=", "2.1")
+    added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
+    controlnet_prompt_embeds = torch.cat((prompt_embeds[1:batch_size], prompt_embeds[1:batch_size]))
+    controlnet_added_cond_kwargs = {key: torch.cat((item[1:batch_size,], item[1:batch_size])) for key, item in added_cond_kwargs.items()}
+    with pipeline.progress_bar(total=num_inference_steps) as progress_bar:
+        for i, t in enumerate(timesteps):
+            # Relevant thread:
+            # https://dev-discuss.pytorch.org/t/cudagraphs-in-pytorch-2-0/1428
+            if (is_unet_compiled and is_controlnet_compiled) and is_torch_higher_equal_2_1:
+                torch._inductor.cudagraph_mark_step_begin()
+            # expand the latents if we are doing classifier free guidance
+            latent_model_input = torch.cat([latents] * 2)
+            latent_model_input = pipeline.scheduler.scale_model_input(latent_model_input, t)           
+
+            # controlnet(s) inference
+            control_model_input = torch.cat((latent_model_input[1:batch_size,], latent_model_input[batch_size+1:]))
+
+            if isinstance(controlnet_keep[i], list):
+                cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])]
+            else:
+                controlnet_cond_scale = controlnet_conditioning_scale
+                if isinstance(controlnet_cond_scale, list):
+                    controlnet_cond_scale = controlnet_cond_scale[0]
+                cond_scale = controlnet_cond_scale * controlnet_keep[i]
+            if cond_scale > 0:
+                down_block_res_samples, mid_block_res_sample = pipeline.controlnet(
+                    control_model_input,
+                    t,
+                    encoder_hidden_states=controlnet_prompt_embeds,
+                    controlnet_cond=image,
+                    conditioning_scale=cond_scale,
+                    guess_mode=False,
+                    added_cond_kwargs=controlnet_added_cond_kwargs,
+                    return_dict=False,
+                )
+    
+                mid_block_res_sample = concat_zero_control(mid_block_res_sample)
+                down_block_res_samples =  [concat_zero_control(down_block_res_sample) for down_block_res_sample in down_block_res_samples]
+            else:
+                mid_block_res_sample = down_block_res_samples = None
+            # predict the noise residual
+            noise_pred = pipeline.unet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=prompt_embeds,
+                timestep_cond=timestep_cond,
+                cross_attention_kwargs=cross_attention_kwargs,
+                down_block_additional_residuals=down_block_res_samples,
+                mid_block_additional_residual=mid_block_res_sample,
+                added_cond_kwargs=added_cond_kwargs,
+                return_dict=False,
+            )[0]
+
+            # perform guidance
+            noise_pred_uncond, noise_pred_text = noise_pred.chunk(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 = pipeline.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % pipeline.scheduler.order == 0):
+                progress_bar.update()
+               
+    # manually for max memory savings
+    if pipeline.vae.dtype == torch.float16 and pipeline.vae.config.force_upcast:
+        pipeline.upcast_vae()
+        latents = latents.to(next(iter(pipeline.vae.post_quant_conv.parameters())).dtype)
+
+    # make sure the VAE is in float32 mode, as it overflows in float16
+    needs_upcasting = pipeline.vae.dtype == torch.float16 and pipeline.vae.config.force_upcast
+
+    if needs_upcasting:
+        pipeline.upcast_vae()
+        latents = latents.to(next(iter(pipeline.vae.post_quant_conv.parameters())).dtype)
+
+    image = pipeline.vae.decode(latents / pipeline.vae.config.scaling_factor, return_dict=False)[0]
+
+    # cast back to fp16 if needed
+    if needs_upcasting:
+        pipeline.vae.to(dtype=torch.float16)
+ 
+    if pipeline.watermark is not None:
+        image = pipeline.watermark.apply_watermark(image)
+
+    image = pipeline.image_processor.postprocess(image, output_type='pil')
+
+    # Offload all models
+    pipeline.maybe_free_model_hooks()
+    return image
+
+
+@torch.no_grad()
+def panorama_call(
+        pipeline: StableDiffusionPanoramaPipeline,
+        prompt: list[str],
+        height: int | None = 512,
+        width: int | None = 2048,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        view_batch_size: int = 1,
+        negative_prompt: str | list[str] | None = None,
+        num_images_per_prompt: int | None = 1,
+        eta: float = 0.0,
+        generator: torch.Generator | None = None,
+        reference_latent: TN = None,
+        latents: TN = None,
+        prompt_embeds: TN = None,
+        negative_prompt_embeds: TN = None,
+        cross_attention_kwargs: dict[str, Any] | None = None,
+        circular_padding: bool = False,
+        clip_skip: int | None = None,
+        stride=8
+) -> list[Image]:
+    # 0. Default height and width to unet
+    height = height or pipeline.unet.config.sample_size * pipeline.vae_scale_factor
+    width = width or pipeline.unet.config.sample_size * pipeline.vae_scale_factor
+
+    # 1. Check inputs. Raise error if not correct
+    pipeline.check_inputs(
+        prompt, height, width, 1, negative_prompt, prompt_embeds, negative_prompt_embeds
+    )
+
+    device = pipeline._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_classifier_free_guidance = guidance_scale > 1.0
+
+    # 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 = pipeline.encode_prompt(
+        prompt,
+        device,
+        num_images_per_prompt,
+        do_classifier_free_guidance,
+        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
+
+    # 4. Prepare timesteps
+    pipeline.scheduler.set_timesteps(num_inference_steps, device=device)
+    timesteps = pipeline.scheduler.timesteps
+
+    # 5. Prepare latent variables
+    num_channels_latents = pipeline.unet.config.in_channels
+    latents = pipeline.prepare_latents(
+        1,
+        num_channels_latents,
+        height,
+        width,
+        prompt_embeds.dtype,
+        device,
+        generator,
+        latents,
+    )
+    if reference_latent is None:
+        reference_latent = torch.randn(1, 4, pipeline.unet.config.sample_size, pipeline.unet.config.sample_size,
+                                       generator=generator)
+    reference_latent = reference_latent.to(device=device, dtype=pipeline.unet.dtype)
+    # 6. Define panorama grid and initialize views for synthesis.
+    # prepare batch grid
+    views = pipeline.get_views(height, width, circular_padding=circular_padding, stride=stride)
+    views_batch = [views[i: i + view_batch_size] for i in range(0, len(views), view_batch_size)]
+    views_scheduler_status = [copy.deepcopy(pipeline.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 = pipeline.prepare_extra_step_kwargs(generator, eta)
+
+    # 8. Denoising loop
+    # Each denoising step also includes refinement of the latents with respect to the
+    # views.
+    num_warmup_steps = len(timesteps) - num_inference_steps * pipeline.scheduler.order
+
+    negative_prompt_embeds = torch.cat([negative_prompt_embeds[:1],
+                                        *[negative_prompt_embeds[1:]] * view_batch_size]
+                                       )
+    prompt_embeds = torch.cat([prompt_embeds[:1],
+                               *[prompt_embeds[1:]] * view_batch_size]
+                              )
+
+    with pipeline.progress_bar(total=num_inference_steps) as progress_bar:
+        for i, t in enumerate(timesteps):
+            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
+                if circular_padding:
+                    latents_for_view = []
+                    for h_start, h_end, w_start, w_end in batch_view:
+                        if w_end > latents.shape[3]:
+                            # Add circular horizontal padding
+                            latent_view = torch.cat(
+                                (
+                                    latents[:, :, h_start:h_end, w_start:],
+                                    latents[:, :, h_start:h_end, : w_end - latents.shape[3]],
+                                ),
+                                dim=-1,
+                            )
+                        else:
+                            latent_view = latents[:, :, h_start:h_end, w_start:w_end]
+                        latents_for_view.append(latent_view)
+                    latents_for_view = torch.cat(latents_for_view)
+                else:
+                    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
+                pipeline.scheduler.__dict__.update(views_scheduler_status[j])
+
+                # expand the latents if we are doing classifier free guidance
+                latent_reference_plus_view = torch.cat((reference_latent, latents_for_view))
+                latent_model_input = latent_reference_plus_view.repeat(2, 1, 1, 1)
+                prompt_embeds_input = torch.cat([negative_prompt_embeds[: 1 + vb_size],
+                                                 prompt_embeds[: 1 + vb_size]]
+                                                )
+                latent_model_input = pipeline.scheduler.scale_model_input(latent_model_input, t)
+                # predict the noise residual
+                # return
+                noise_pred = pipeline.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds_input,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                ).sample
+
+                # perform guidance
+
+                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+                # compute the previous noisy sample x_t -> x_t-1
+                latent_reference_plus_view = pipeline.scheduler.step(
+                    noise_pred, t, latent_reference_plus_view, **extra_step_kwargs
+                ).prev_sample
+                if j == len(views_batch) - 1:
+                    reference_latent = latent_reference_plus_view[:1]
+                latents_denoised_batch = latent_reference_plus_view[1:]
+                # save views scheduler status after sample
+                views_scheduler_status[j] = copy.deepcopy(pipeline.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
+                ):
+                    if circular_padding and w_end > latents.shape[3]:
+                        # Case for circular padding
+                        value[:, :, h_start:h_end, w_start:] += latents_view_denoised[
+                                                                :, :, h_start:h_end, : latents.shape[3] - w_start
+                                                                ]
+                        value[:, :, h_start:h_end, : w_end - latents.shape[3]] += latents_view_denoised[
+                                                                                  :, :, h_start:h_end,
+                                                                                  latents.shape[3] - w_start:
+                                                                                  ]
+                        count[:, :, h_start:h_end, w_start:] += 1
+                        count[:, :, h_start:h_end, : w_end - latents.shape[3]] += 1
+                    else:
+                        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)
+
+            # call the callback, if provided
+            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % pipeline.scheduler.order == 0):
+                progress_bar.update()
+
+    if circular_padding:
+        image = pipeline.decode_latents_with_padding(latents)
+    else:
+        image = pipeline.vae.decode(latents / pipeline.vae.config.scaling_factor, return_dict=False)[0]
+    reference_image = pipeline.vae.decode(reference_latent / pipeline.vae.config.scaling_factor, return_dict=False)[0]
+    # image, has_nsfw_concept = pipeline.run_safety_checker(image, device, prompt_embeds.dtype)
+    # reference_image, _ = pipeline.run_safety_checker(reference_image, device, prompt_embeds.dtype)
+
+    image = pipeline.image_processor.postprocess(image, output_type='pil', do_denormalize=[True])
+    reference_image = pipeline.image_processor.postprocess(reference_image, output_type='pil', do_denormalize=[True])
+    pipeline.maybe_free_model_hooks()
+    return reference_image + image

requirements.txt

@@ -1,4 +1,5 @@
-diffusers
-accelerate
+diffusers==0.16.1
+transformers
 mediapy
+ipywidgets
 einops

sa_handler.py

@@ -0,0 +1,269 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from diffusers import StableDiffusionXLPipeline
+import torch
+import torch.nn as nn
+from torch.nn import functional as nnf
+from diffusers.models import attention_processor
+import einops
+
+T = torch.Tensor
+
+
+@dataclass(frozen=True)
+class StyleAlignedArgs:
+    share_group_norm: bool = True
+    share_layer_norm: bool = True,
+    share_attention: bool = True
+    adain_queries: bool = True
+    adain_keys: bool = True
+    adain_values: bool = False
+    full_attention_share: bool = False
+    keys_scale: float = 1.
+    only_self_level: float = 0.
+
+
+def expand_first(feat: T, scale=1., ) -> T:
+    b = feat.shape[0]
+    feat_style = torch.stack((feat[0], feat[b // 2])).unsqueeze(1)
+    if scale == 1:
+        feat_style = feat_style.expand(2, b // 2, *feat.shape[1:])
+    else:
+        feat_style = feat_style.repeat(1, b // 2, 1, 1, 1)
+        feat_style = torch.cat([feat_style[:, :1], scale * feat_style[:, 1:]], dim=1)
+    return feat_style.reshape(*feat.shape)
+
+
+def concat_first(feat: T, dim=2, scale=1.) -> T:
+    feat_style = expand_first(feat, scale=scale)
+    return torch.cat((feat, feat_style), dim=dim)
+
+
+def calc_mean_std(feat, eps: float = 1e-5) -> tuple[T, T]:
+    feat_std = (feat.var(dim=-2, keepdims=True) + eps).sqrt()
+    feat_mean = feat.mean(dim=-2, keepdims=True)
+    return feat_mean, feat_std
+
+
+def adain(feat: T) -> T:
+    feat_mean, feat_std = calc_mean_std(feat)
+    feat_style_mean = expand_first(feat_mean)
+    feat_style_std = expand_first(feat_std)
+    feat = (feat - feat_mean) / feat_std
+    feat = feat * feat_style_std + feat_style_mean
+    return feat
+
+
+class DefaultAttentionProcessor(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.processor = attention_processor.AttnProcessor2_0()
+
+    def __call__(self, attn: attention_processor.Attention, hidden_states, encoder_hidden_states=None,
+                 attention_mask=None, **kwargs):
+        return self.processor(attn, hidden_states, encoder_hidden_states, attention_mask)
+
+
+class SharedAttentionProcessor(DefaultAttentionProcessor):
+
+    def shared_call(
+            self,
+            attn: attention_processor.Attention,
+            hidden_states,
+            encoder_hidden_states=None,
+            attention_mask=None,
+            **kwargs
+    ):
+
+        residual = hidden_states
+        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 = attn.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, attn.heads, -1, attention_mask.shape[-1])
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        # if self.step >= self.start_inject:
+        if self.adain_queries:
+            query = adain(query)
+        if self.adain_keys:
+            key = adain(key)
+        if self.adain_values:
+            value = adain(value)
+        if self.share_attention:
+            key = concat_first(key, -2, scale=self.keys_scale)
+            value = concat_first(value, -2)
+            hidden_states = nnf.scaled_dot_product_attention(
+                query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+            )
+        else:
+            hidden_states = nnf.scaled_dot_product_attention(
+                query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+            )
+        # hidden_states = adain(hidden_states)
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+        return hidden_states
+
+    def __call__(self, attn: attention_processor.Attention, hidden_states, encoder_hidden_states=None,
+                 attention_mask=None, **kwargs):
+        if self.full_attention_share:
+            b, n, d = hidden_states.shape
+            hidden_states = einops.rearrange(hidden_states, '(k b) n d -> k (b n) d', k=2)
+            hidden_states = super().__call__(attn, hidden_states, encoder_hidden_states=encoder_hidden_states,
+                                             attention_mask=attention_mask, **kwargs)
+            hidden_states = einops.rearrange(hidden_states, 'k (b n) d -> (k b) n d', n=n)
+        else:
+            hidden_states = self.shared_call(attn, hidden_states, hidden_states, attention_mask, **kwargs)
+
+        return hidden_states
+
+    def __init__(self, style_aligned_args: StyleAlignedArgs):
+        super().__init__()
+        self.share_attention = style_aligned_args.share_attention
+        self.adain_queries = style_aligned_args.adain_queries
+        self.adain_keys = style_aligned_args.adain_keys
+        self.adain_values = style_aligned_args.adain_values
+        self.full_attention_share = style_aligned_args.full_attention_share
+        self.keys_scale = style_aligned_args.keys_scale
+
+
+def _get_switch_vec(total_num_layers, level):
+    if level == 0:
+        return torch.zeros(total_num_layers, dtype=torch.bool)
+    if level == 1:
+        return torch.ones(total_num_layers, dtype=torch.bool)
+    to_flip = level > .5
+    if to_flip:
+        level = 1 - level
+    num_switch = int(level * total_num_layers)
+    vec = torch.arange(total_num_layers)
+    vec = vec % (total_num_layers // num_switch)
+    vec = vec == 0
+    if to_flip:
+        vec = ~vec
+    return vec
+
+
+def init_attention_processors(pipeline: StableDiffusionXLPipeline, style_aligned_args: StyleAlignedArgs | None = None):
+    attn_procs = {}
+    unet = pipeline.unet
+    number_of_self, number_of_cross = 0, 0
+    num_self_layers = len([name for name in unet.attn_processors.keys() if 'attn1' in name])
+    if style_aligned_args is None:
+        only_self_vec = _get_switch_vec(num_self_layers, 1)
+    else:
+        only_self_vec = _get_switch_vec(num_self_layers, style_aligned_args.only_self_level)
+    for i, name in enumerate(unet.attn_processors.keys()):
+        is_self_attention = 'attn1' in name
+        if is_self_attention:
+            number_of_self += 1
+            if style_aligned_args is None or only_self_vec[i // 2]:
+                attn_procs[name] = DefaultAttentionProcessor()
+            else:
+                attn_procs[name] = SharedAttentionProcessor(style_aligned_args)
+
+        else:
+            number_of_cross += 1
+            attn_procs[name] = DefaultAttentionProcessor()
+
+    unet.set_attn_processor(attn_procs)
+
+
+def register_shared_norm(pipeline: StableDiffusionXLPipeline,
+                         share_group_norm: bool = True,
+                         share_layer_norm: bool = True, ):
+    def register_norm_forward(norm_layer: nn.GroupNorm | nn.LayerNorm) -> nn.GroupNorm | nn.LayerNorm:
+        if not hasattr(norm_layer, 'orig_forward'):
+            setattr(norm_layer, 'orig_forward', norm_layer.forward)
+        orig_forward = norm_layer.orig_forward
+
+        def forward_(hidden_states: T) -> T:
+            n = hidden_states.shape[-2]
+            hidden_states = concat_first(hidden_states, dim=-2)
+            hidden_states = orig_forward(hidden_states)
+            return hidden_states[..., :n, :]
+
+        norm_layer.forward = forward_
+        return norm_layer
+
+    def get_norm_layers(pipeline_, norm_layers_: dict[str, list[nn.GroupNorm | nn.LayerNorm]]):
+        if isinstance(pipeline_, nn.LayerNorm) and share_layer_norm:
+            norm_layers_['layer'].append(pipeline_)
+        if isinstance(pipeline_, nn.GroupNorm) and share_group_norm:
+            norm_layers_['group'].append(pipeline_)
+        else:
+            for layer in pipeline_.children():
+                get_norm_layers(layer, norm_layers_)
+
+    norm_layers = {'group': [], 'layer': []}
+    get_norm_layers(pipeline.unet, norm_layers)
+    return [register_norm_forward(layer) for layer in norm_layers['group']] + [register_norm_forward(layer) for layer in
+                                                                               norm_layers['layer']]
+
+
+class Handler:
+
+    def register(self, style_aligned_args: StyleAlignedArgs, ):
+        self.norm_layers = register_shared_norm(self.pipeline, style_aligned_args.share_group_norm,
+                                                style_aligned_args.share_layer_norm)
+        init_attention_processors(self.pipeline, style_aligned_args)
+
+    def remove(self):
+        for layer in self.norm_layers:
+            layer.forward = layer.orig_forward
+        self.norm_layers = []
+        init_attention_processors(self.pipeline, None)
+
+    def __init__(self, pipeline: StableDiffusionXLPipeline):
+        self.pipeline = pipeline
+        self.norm_layers = []

style_aligned_sdxl.ipynb

@@ -0,0 +1,142 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "a885cf5d-c525-4f5b-a8e4-f67d2f699909",
+   "metadata": {},
+   "source": [
+    "## Copyright 2023 Google LLC"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d891d022-8979-40d4-848f-ecb84c17f12c",
+   "metadata": {
+    "jp-MarkdownHeadingCollapsed": true
+   },
+   "outputs": [],
+   "source": [
+    "# Copyright 2023 Google LLC\n",
+    "#\n",
+    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+    "# you may not use this file except in compliance with the License.\n",
+    "# You may obtain a copy of the License at\n",
+    "#\n",
+    "#      http://www.apache.org/licenses/LICENSE-2.0\n",
+    "#\n",
+    "# Unless required by applicable law or agreed to in writing, software\n",
+    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+    "# See the License for the specific language governing permissions and\n",
+    "# limitations under the License."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "540d8642-c203-471c-a66d-0d43aabb0706",
+   "metadata": {},
+   "source": [
+    "# StyleAligned over SDXL"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "23d54ea7-f7ab-4548-9b10-ece87216dc18",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from diffusers import StableDiffusionXLPipeline, DDIMScheduler\n",
+    "import torch\n",
+    "import mediapy\n",
+    "import sa_handler"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c2f6f1e6-445f-47bc-b9db-0301caeb7490",
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# init models\n",
+    "\n",
+    "scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule=\"scaled_linear\", clip_sample=False,\n",
+    "                              set_alpha_to_one=False)\n",
+    "pipeline = StableDiffusionXLPipeline.from_pretrained(\n",
+    "    \"stabilityai/stable-diffusion-xl-base-1.0\", torch_dtype=torch.float16, variant=\"fp16\", use_safetensors=True,\n",
+    "    scheduler=scheduler\n",
+    ").to(\"cuda\")\n",
+    "\n",
+    "handler = sa_handler.Handler(pipeline)\n",
+    "sa_args = sa_handler.StyleAlignedArgs(share_group_norm=False,\n",
+    "                                      share_layer_norm=False,\n",
+    "                                      share_attention=True,\n",
+    "                                      adain_queries=True,\n",
+    "                                      adain_keys=True,\n",
+    "                                      adain_values=False,\n",
+    "                                     )\n",
+    "\n",
+    "handler.register(sa_args, )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5cca9256-0ce0-45c3-9cba-68c7eff1452f",
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# run StyleAligned\n",
+    "\n",
+    "sets_of_prompts = [\n",
+    "  \"a toy train. macro photo. 3d game asset\",\n",
+    "  \"a toy airplane. macro photo. 3d game asset\",\n",
+    "  \"a toy bicycle. macro photo. 3d game asset\",\n",
+    "  \"a toy car. macro photo. 3d game asset\",\n",
+    "  \"a toy boat. macro photo. 3d game asset\",\n",
+    "]\n",
+    "images = pipeline(sets_of_prompts,).images\n",
+    "mediapy.show_images(images)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d819ad6d-0c19-411f-ba97-199909f64805",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

style_aligned_w_controlnet.ipynb

@@ -0,0 +1,200 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "f86ede39-8d9f-4da9-bc12-955f2fddd484",
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "## Copyright 2023 Google LLC"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3f3cbf47-a52b-48b1-9bd3-3435f92f2174",
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# Copyright 2023 Google LLC\n",
+    "#\n",
+    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+    "# you may not use this file except in compliance with the License.\n",
+    "# You may obtain a copy of the License at\n",
+    "#\n",
+    "#      http://www.apache.org/licenses/LICENSE-2.0\n",
+    "#\n",
+    "# Unless required by applicable law or agreed to in writing, software\n",
+    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+    "# See the License for the specific language governing permissions and\n",
+    "# limitations under the License."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "22de629b-581f-4335-9e7b-f73221d8dbcb",
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "# ControlNet depth with StyleAligned over SDXL"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "486b7ebb-c483-4bf0-ace8-f8092c2d1f23",
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL\n",
+    "from diffusers.utils import load_image\n",
+    "from transformers import DPTImageProcessor, DPTForDepthEstimation\n",
+    "import torch\n",
+    "import mediapy\n",
+    "import sa_handler\n",
+    "import pipeline_calls"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2a7e85e7-b5cf-45b2-946a-5ba1e4923586",
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# init models\n",
+    "\n",
+    "depth_estimator = DPTForDepthEstimation.from_pretrained(\"Intel/dpt-hybrid-midas\").to(\"cuda\")\n",
+    "feature_processor = DPTImageProcessor.from_pretrained(\"Intel/dpt-hybrid-midas\")\n",
+    "\n",
+    "controlnet = ControlNetModel.from_pretrained(\n",
+    "    \"diffusers/controlnet-depth-sdxl-1.0\",\n",
+    "    variant=\"fp16\",\n",
+    "    use_safetensors=True,\n",
+    "    torch_dtype=torch.float16,\n",
+    ").to(\"cuda\")\n",
+    "vae = AutoencoderKL.from_pretrained(\"madebyollin/sdxl-vae-fp16-fix\", torch_dtype=torch.float16).to(\"cuda\")\n",
+    "pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(\n",
+    "    \"stabilityai/stable-diffusion-xl-base-1.0\",\n",
+    "    controlnet=controlnet,\n",
+    "    vae=vae,\n",
+    "    variant=\"fp16\",\n",
+    "    use_safetensors=True,\n",
+    "    torch_dtype=torch.float16,\n",
+    ").to(\"cuda\")\n",
+    "pipeline.enable_model_cpu_offload()\n",
+    "\n",
+    "sa_args = sa_handler.StyleAlignedArgs(share_group_norm=False,\n",
+    "                                      share_layer_norm=False,\n",
+    "                                      share_attention=True,\n",
+    "                                      adain_queries=True,\n",
+    "                                      adain_keys=True,\n",
+    "                                      adain_values=False,\n",
+    "                                     )\n",
+    "handler = sa_handler.Handler(pipeline)\n",
+    "handler.register(sa_args, )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "94ca26b4-9061-4012-9400-8d97ef212d87",
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# get depth maps\n",
+    "\n",
+    "image = load_image(\"./example_image/train.png\")\n",
+    "depth_image1 = pipeline_calls.get_depth_map(image, feature_processor, depth_estimator)\n",
+    "depth_image2 = load_image(\"./example_image/sun.png\").resize((1024, 1024))\n",
+    "mediapy.show_images([depth_image1, depth_image2])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c8f56fe4-559f-49ff-a2d8-460dcfeb56a0",
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# run ControlNet depth with StyleAligned\n",
+    "\n",
+    "reference_prompt = \"a poster in flat design style\"\n",
+    "target_prompts = [\"a train in flat design style\", \"the sun in flat design style\"]\n",
+    "controlnet_conditioning_scale = 0.8\n",
+    "num_images_per_prompt = 3 # adjust according to VRAM size\n",
+    "latents = torch.randn(1 + num_images_per_prompt, 4, 128, 128).to(pipeline.unet.dtype)\n",
+    "for deph_map, target_prompt in zip((depth_image1, depth_image2), target_prompts):\n",
+    "    latents[1:] = torch.randn(num_images_per_prompt, 4, 128, 128).to(pipeline.unet.dtype)\n",
+    "    images = pipeline_calls.controlnet_call(pipeline, [reference_prompt, target_prompt],\n",
+    "                                            image=deph_map,\n",
+    "                                            num_inference_steps=50,\n",
+    "                                            controlnet_conditioning_scale=controlnet_conditioning_scale,\n",
+    "                                            num_images_per_prompt=num_images_per_prompt,\n",
+    "                                           latents=latents)\n",
+    "    \n",
+    "    mediapy.show_images([images[0], deph_map] +  images[1:], titles=[\"reference\", \"depth\"] + [f'result {i}' for i in range(1, len(images))])\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "437ba4bd-6243-486b-8ba5-3b7cd661d53a",
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

style_aligned_w_multidiffusion.ipynb

@@ -0,0 +1,156 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "50fa980f-1bae-40c1-a1f3-f5f89bef60d3",
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "## Copyright 2023 Google LLC"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5da5f038-057f-4475-a783-95660f98238c",
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# Copyright 2023 Google LLC\n",
+    "#\n",
+    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+    "# you may not use this file except in compliance with the License.\n",
+    "# You may obtain a copy of the License at\n",
+    "#\n",
+    "#      http://www.apache.org/licenses/LICENSE-2.0\n",
+    "#\n",
+    "# Unless required by applicable law or agreed to in writing, software\n",
+    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+    "# See the License for the specific language governing permissions and\n",
+    "# limitations under the License."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c3a7c069-c441-4204-a905-59cbd9edc13a",
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "# MultiDiffusion with StyleAligned over SD v2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "14178de7-d4c8-4881-ac1d-ff84bae57c6f",
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from diffusers import StableDiffusionPanoramaPipeline, DDIMScheduler\n",
+    "import mediapy\n",
+    "import sa_handler\n",
+    "import pipeline_calls"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "738cee0e-4d6e-4875-b4df-eadff6e27e7f",
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# init models\n",
+    "model_ckpt = \"stabilityai/stable-diffusion-2-base\"\n",
+    "scheduler = DDIMScheduler.from_pretrained(model_ckpt, subfolder=\"scheduler\")\n",
+    "pipeline = StableDiffusionPanoramaPipeline.from_pretrained(\n",
+    "     model_ckpt, scheduler=scheduler, torch_dtype=torch.float16\n",
+    ").to(\"cuda\")\n",
+    "\n",
+    "sa_args = sa_handler.StyleAlignedArgs(share_group_norm=True,\n",
+    "                                      share_layer_norm=True,\n",
+    "                                      share_attention=True,\n",
+    "                                      adain_queries=True,\n",
+    "                                      adain_keys=True,\n",
+    "                                      adain_values=False,\n",
+    "                                     )\n",
+    "handler = sa_handler.Handler(pipeline)\n",
+    "handler.register(sa_args)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ea61e789-2814-4820-8ae7-234c3c6640a0",
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# run MultiDiffusion with StyleAligned\n",
+    "\n",
+    "reference_prompt = \"a beautiful papercut art design\"\n",
+    "target_prompts = [\"mountains in a beautiful papercut art design\", \"giraffes in a beautiful papercut art design\"]\n",
+    "view_batch_size = 25  # adjust according to VRAM size\n",
+    "reference_latent = torch.randn(1, 4, 64, 64,)\n",
+    "for target_prompt in target_prompts:\n",
+    "    images = pipeline_calls.panorama_call(pipeline, [reference_prompt, target_prompt], reference_latent=reference_latent, view_batch_size=view_batch_size)\n",
+    "    mediapy.show_images(images, titles=[\"reference\", \"result\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "791a9b28-f0ce-4fd0-9f3c-594281c2ae56",
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}