ZeroGPU support (#2)

- Support ZeroGPU (aca4f93b09d9e8b8f7406e9c8579c1943d09ebe3)


Co-authored-by: Apolinário from multimodal AI art <multimodalart@users.noreply.huggingface.co>

.gitattributes

@@ -33,3 +33,7 @@ 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/camel.png 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

@@ -3,7 +3,7 @@ title: StyleAligned Transfer
 emoji: 🐠
 colorFrom: blue
 colorTo: pink
-sdk: docker
+sdk: gradio
 pinned: false
 ---
 

app.py

@@ -6,6 +6,7 @@ import math
 from diffusers.utils import load_image
 import inversion
 import numpy as np
+import spaces
 
 # init models
 
@@ -22,6 +23,7 @@ pipeline = StableDiffusionXLPipeline.from_pretrained(
 pipeline.enable_model_cpu_offload() 
 pipeline.enable_vae_slicing()
 
+@spaces.GPU(duration=120)
 def run(ref_path, ref_style, ref_prompt, prompt1, prompt2, prompt3):
     # DDIM inversion
     src_style = f"{ref_style}"
@@ -41,8 +43,6 @@ def run(ref_path, ref_style, ref_prompt, prompt1, prompt2, prompt3):
     prompts = [
         src_prompt,
         prompt1,
-        prompt2,
-        prompt3
     ]
 
     # some parameters you can adjust to control fidelity to reference
@@ -105,8 +105,8 @@ with gr.Blocks(css=css) as demo:
                 with gr.Group():
                     results = gr.Gallery()
                     prompt1 = gr.Textbox(label="Prompt1", value="A man working on a laptop")
-                    prompt2 = gr.Textbox(label="Prompt2", value="A man eating pizza")
-                    prompt3 = gr.Textbox(label="Prompt3", value="A woman playing on saxophone")
+                    prompt2 = gr.Textbox(label="Prompt2", value="A man eating pizza", visible=False)
+                    prompt3 = gr.Textbox(label="Prompt3", value="A woman playing on saxophone", visible=False)
                     run_button = gr.Button("Submit")
             
                 

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/).

demo_stylealigned_controlnet.py

@@ -0,0 +1,139 @@
+import gradio as gr
+from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL
+from diffusers.utils import load_image
+from transformers import DPTImageProcessor, DPTForDepthEstimation
+import torch
+import sa_handler
+import pipeline_calls
+
+
+
+# Initialize models
+depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to("cuda")
+feature_processor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas")
+
+controlnet = ControlNetModel.from_pretrained(
+    "diffusers/controlnet-depth-sdxl-1.0",
+    variant="fp16",
+    use_safetensors=True,
+    torch_dtype=torch.float16,
+).to("cuda")
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16).to("cuda")
+pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    controlnet=controlnet,
+    vae=vae,
+    variant="fp16",
+    use_safetensors=True,
+    torch_dtype=torch.float16,
+).to("cuda")
+# Configure pipeline for CPU offloading and VAE slicing
+pipeline.enable_model_cpu_offload()
+pipeline.enable_vae_slicing()
+
+# Initialize style-aligned handler
+sa_args = sa_handler.StyleAlignedArgs(share_group_norm=False,
+                                      share_layer_norm=False,
+                                      share_attention=True,
+                                      adain_queries=True,
+                                      adain_keys=True,
+                                      adain_values=False,
+                                     )
+handler = sa_handler.Handler(pipeline)
+handler.register(sa_args, )
+
+
+# Function to run ControlNet depth with StyleAligned
+def style_aligned_controlnet(ref_style_prompt, depth_map, ref_image, img_generation_prompt, seed):
+    try:
+        if depth_map == True:
+            image = load_image(ref_image)
+            depth_image = pipeline_calls.get_depth_map(image, feature_processor, depth_estimator)
+        else:
+            depth_image = load_image(ref_image).resize((1024, 1024))
+        controlnet_conditioning_scale = 0.8
+        gen = None if seed is None else torch.manual_seed(int(seed))
+        num_images_per_prompt = 3  # adjust according to VRAM size
+        latents = torch.randn(1 + num_images_per_prompt, 4, 128, 128, generator=gen).to(pipeline.unet.dtype)
+
+        images = pipeline_calls.controlnet_call(pipeline, [ref_style_prompt, img_generation_prompt],
+                                                image=depth_image,
+                                                num_inference_steps=50,
+                                                controlnet_conditioning_scale=controlnet_conditioning_scale,
+                                                num_images_per_prompt=num_images_per_prompt,
+                                                latents=latents)
+        return [images[0], depth_image] + images[1:], gr.Image(value=images[0], visible=True)
+    except Exception as e:
+        raise gr.Error(f"Error in generating images:{e}")
+
+# Create a Gradio UI
+with gr.Blocks() as demo:
+    gr.HTML('<h1 style="text-align: center;">ControlNet with StyleAligned</h1>')
+    with gr.Row():
+      
+      with gr.Column(variant='panel'):
+        # Textbox for reference style prompt
+        ref_style_prompt = gr.Textbox(
+          label='Reference style prompt',
+          info="Enter a Prompt to generate the reference image", placeholder='a poster in <style name> style'
+        )
+        with gr.Row(variant='panel'):
+        # Checkbox for using controller depth-map
+            depth_map = gr.Checkbox(label='Depth-map',)
+            seed = gr.Number(value=1234, label="Seed", precision=0, step=1, scale=3,
+                                    info="Enter a seed of a previous reference image "
+                                         "or leave empty for a random generation.")
+        # Image display for the generated reference style image
+        ref_style_image = gr.Image(visible=False, label='Reference style image', scale=1)
+
+
+      with gr.Column(variant='panel'): 
+        # Image upload option for uploading a reference image for controlnet
+        ref_image = gr.Image(label="Upload the reference image", 
+                             type='filepath' )
+        # Textbox for ControlNet prompt
+        img_generation_prompt = gr.Textbox(
+            label='Generation Prompt',
+            info="Enter a Prompt to generate images using ControlNet and StyleAligned",
+            )
+
+    # Button to trigger image generation
+    btn = gr.Button("Generate", size='sm')
+    # Gallery to display generated images
+    gallery = gr.Gallery(label="Style-Aligned ControlNet - Generated images", 
+                           elem_id="gallery",
+                           columns=5, 
+                           rows=1, 
+                           object_fit="contain", 
+                           height="auto",
+                          )
+      
+    btn.click(fn=style_aligned_controlnet, 
+              inputs=[ref_style_prompt, depth_map, ref_image, img_generation_prompt, seed],
+              outputs=[gallery, ref_style_image], 
+              api_name="style_aligned_controlnet")
+
+
+    # Example inputs for the Gradio interface
+    gr.Examples(
+      examples=[
+        ['A couple sitting a wooden bench, in colorful clay animation, claymation style.', True,
+         'example_image/train.png', 'A train in colorful clay animation, claymation style.',],
+        ['A couple sitting a wooden bench, in colorful clay animation, claymation style.', False,
+         'example_image/sun.png', 'Sun in colorful clay animation, claymation style.',],
+        ['A poster in a papercut art style.', False,
+         'example_image/A.png', 'Letter A in a papercut art style.', None],
+        ['A bull in a low-poly, colorful origami style.', True, 'example_image/whale.png',
+         'A whale in a low-poly, colorful origami style.', None],
+        ['An image in ancient egyptian art style, hieroglyphics style.', True, 'example_image/camel.png',
+         'A camel in a painterly, digital illustration style.',],
+        ['An image in ancient egyptian art style, hieroglyphics style.', True, 'example_image/whale.png',
+         'A whale in ancient egyptian art style, hieroglyphics style.',],
+      ],
+      inputs=[ref_style_prompt, depth_map, ref_image, img_generation_prompt,],
+      outputs=[gallery, ref_style_image], 
+      fn=style_aligned_controlnet,
+      )
+
+# Launch the Gradio demo   
+demo.launch()

demo_stylealigned_multidiffusion.py

@@ -0,0 +1,104 @@
+import gradio as gr
+import torch
+from diffusers import StableDiffusionPanoramaPipeline, DDIMScheduler
+import sa_handler
+import pipeline_calls
+
+
+# init models
+model_ckpt = "stabilityai/stable-diffusion-2-base"
+scheduler = DDIMScheduler.from_pretrained(model_ckpt, subfolder="scheduler")
+pipeline = StableDiffusionPanoramaPipeline.from_pretrained(
+     model_ckpt, scheduler=scheduler, torch_dtype=torch.float16
+).to("cuda")
+# Configure the pipeline for CPU offloading and VAE slicing
+pipeline.enable_model_cpu_offload()
+pipeline.enable_vae_slicing()
+sa_args = sa_handler.StyleAlignedArgs(share_group_norm=True,
+                                      share_layer_norm=True,
+                                      share_attention=True,
+                                      adain_queries=True,
+                                      adain_keys=True,
+                                      adain_values=False,
+                                     )
+# Initialize the style-aligned handler
+handler = sa_handler.Handler(pipeline)
+handler.register(sa_args)
+
+
+# Define the function to run MultiDiffusion with StyleAligned
+def style_aligned_multidiff(ref_style_prompt, img_generation_prompt, seed):
+    try:
+        view_batch_size = 25  # adjust according to VRAM size
+        gen = None if seed is None else torch.manual_seed(int(seed))
+        reference_latent = torch.randn(1, 4, 64, 64, generator=gen)
+        images = pipeline_calls.panorama_call(pipeline,
+                                              [ref_style_prompt, img_generation_prompt],
+                                              reference_latent=reference_latent,
+                                              view_batch_size=view_batch_size)
+    
+        return images, gr.Image(value=images[0], visible=True)
+    except Exception as e:
+        raise gr.Error(f"Error in generating images:{e}")
+
+# Create a Gradio UI
+with gr.Blocks() as demo:
+    gr.HTML('<h1 style="text-align: center;">MultiDiffusion with StyleAligned </h1>')
+    with gr.Row():
+      with gr.Column(variant='panel'):
+        # Textbox for reference style prompt
+        ref_style_prompt = gr.Textbox(
+          label='Reference style prompt',
+          info='Enter a Prompt to generate the reference image',
+          placeholder='A poster in a papercut art style.'
+        )
+        seed = gr.Number(value=1234, label="Seed", precision=0, step=1,
+                         info="Enter a seed of a previous reference image "
+                              "or leave empty for a random generation.")
+        # Image display for the reference style image
+        ref_style_image = gr.Image(visible=False, label='Reference style image')
+
+
+      with gr.Column(variant='panel'):
+        # Textbox for prompt for MultiDiffusion panoramas
+        img_generation_prompt = gr.Textbox(
+          label='MultiDiffusion Prompt',
+          info='Enter a Prompt to generate panoramic images using Style-aligned combined with MultiDiffusion',
+          placeholder= 'A village in a papercut art style.'
+          )
+
+    # Button to trigger image generation
+    btn = gr.Button('Style Aligned MultiDiffusion - Generate', size='sm')
+    # Gallery to display generated style image and the panorama
+    gallery = gr.Gallery(label='StyleAligned MultiDiffusion - generated images',
+                           elem_id='gallery',
+                           columns=5,
+                           rows=1,
+                           object_fit='contain',
+                           height='auto',
+                           allow_preview=True,
+                           preview=True,
+                          )
+    # Button click event
+    btn.click(fn=style_aligned_multidiff,
+              inputs=[ref_style_prompt, img_generation_prompt, seed],
+              outputs=[gallery, ref_style_image,],
+              api_name='style_aligned_multidiffusion')
+
+    # Example inputs for the Gradio demo
+    gr.Examples(
+      examples=[
+        ['A poster in a papercut art style.', 'A village in a papercut art style.'],
+        ['A poster in a papercut art style.', 'Futuristic cityscape in a papercut art style.'],
+        ['A poster in a papercut art style.', 'A jungle in a papercut art style.'],
+        ['A poster in a flat design style.', 'Giraffes in a flat design style.'],
+        ['A poster in a flat design style.', 'Houses in a flat design style.'],
+        ['A poster in a flat design style.', 'Mountains in a flat design style.'],
+      ],
+      inputs=[ref_style_prompt, img_generation_prompt],
+      outputs=[gallery, ref_style_image],
+      fn=style_aligned_multidiff,
+      )
+
+# Launch the Gradio demo
+demo.launch()

demo_stylealigned_sdxl.py

@@ -0,0 +1,85 @@
+import gradio as gr
+from diffusers import StableDiffusionXLPipeline, DDIMScheduler
+import torch
+import sa_handler
+
+# init models
+scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False,
+                              set_alpha_to_one=False)
+pipeline = StableDiffusionXLPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True,
+    scheduler=scheduler
+).to("cuda")
+# Configure the pipeline for CPU offloading and VAE slicing#pipeline.enable_sequential_cpu_offload()
+pipeline.enable_model_cpu_offload() 
+pipeline.enable_vae_slicing()
+# Initialize the style-aligned handler
+handler = sa_handler.Handler(pipeline)
+sa_args = sa_handler.StyleAlignedArgs(share_group_norm=False,
+                                      share_layer_norm=False,
+                                      share_attention=True,
+                                      adain_queries=True,
+                                      adain_keys=True,
+                                      adain_values=False,
+                                     )
+
+handler.register(sa_args, )
+
+# Define the function to generate style-aligned images
+def style_aligned_sdxl(initial_prompt1, initial_prompt2, initial_prompt3, initial_prompt4,
+                       initial_prompt5, style_prompt, seed):
+    try:
+        # Combine the style prompt with each initial prompt
+        gen = None if seed is None else torch.manual_seed(int(seed))
+        sets_of_prompts = [prompt + " in the style of " + style_prompt for prompt in [initial_prompt1, initial_prompt2, initial_prompt3, initial_prompt4, initial_prompt5] if prompt]
+        # Generate images using the pipeline
+        images = pipeline(sets_of_prompts, generator=gen).images
+        return images
+    except Exception as e:
+        raise gr.Error(f"Error in generating images: {e}")
+
+with gr.Blocks() as demo:
+    gr.HTML('<h1 style="text-align: center;">StyleAligned SDXL</h1>')
+    with gr.Group():
+      with gr.Column():
+        with gr.Accordion(label='Enter upto 5 different initial prompts', open=True):
+          with gr.Row(variant='panel'):
+            # Textboxes for initial prompts
+            initial_prompt1 = gr.Textbox(label='Initial prompt 1', value='', show_label=False, container=False, placeholder='a toy train')
+            initial_prompt2 = gr.Textbox(label='Initial prompt 2', value='', show_label=False, container=False, placeholder='a toy airplane')
+            initial_prompt3 = gr.Textbox(label='Initial prompt 3', value='', show_label=False, container=False, placeholder='a toy bicycle')
+            initial_prompt4 = gr.Textbox(label='Initial prompt 4', value='', show_label=False, container=False, placeholder='a toy car')
+            initial_prompt5 = gr.Textbox(label='Initial prompt 5', value='', show_label=False, container=False, placeholder='a toy boat')
+        with gr.Row():
+          # Textbox for the style prompt
+          style_prompt = gr.Textbox(label="Enter a style prompt", placeholder='macro photo, 3d game asset', scale=3)
+          seed = gr.Number(value=1234, label="Seed", precision=0, step=1, scale=1,
+                           info="Enter a seed of a previous run "
+                                "or leave empty for a random generation.")
+        # Button to generate images
+        btn = gr.Button("Generate a set of Style-aligned SDXL images",)
+    # Display the generated images
+    output = gr.Gallery(label="Style aligned text-to-image on SDXL ", elem_id="gallery",columns=5, rows=1,
+                        object_fit="contain", height="auto",)
+
+    # Button click event
+    btn.click(fn=style_aligned_sdxl, 
+              inputs=[initial_prompt1, initial_prompt2, initial_prompt3, initial_prompt4, initial_prompt5,
+                      style_prompt, seed],
+              outputs=output, 
+              api_name="style_aligned_sdxl")
+
+    # Providing Example inputs for the demo
+    gr.Examples(examples=[
+                    ["a toy train", "a toy airplane", "a toy bicycle", "a toy car", "a toy boat", "macro photo. 3d game asset."],
+                    ["a toy train", "a toy airplane", "a toy bicycle", "a toy car", "a toy boat", "BW logo. high contrast."],
+                    ["a cat", "a dog", "a bear", "a man on a bicycle", "a girl working on laptop", "minimal origami."],
+                    ["a firewoman", "a Gardner", "a scientist", "a policewoman", "a saxophone player", "made of claymation, stop motion animation."],
+                    ["a firewoman", "a Gardner", "a scientist", "a policewoman", "a saxophone player", "sketch, character sheet."],
+                    ],
+            inputs=[initial_prompt1, initial_prompt2, initial_prompt3, initial_prompt4, initial_prompt5, style_prompt],
+            outputs=[output],
+            fn=style_aligned_sdxl)
+
+# Launch the Gradio demo
+demo.launch()

inversion.py

@@ -0,0 +1,125 @@
+# 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 Callable
+from diffusers import StableDiffusionXLPipeline
+import torch
+from tqdm import tqdm
+import numpy as np
+
+
+T = torch.Tensor
+TN = T | None
+InversionCallback = Callable[[StableDiffusionXLPipeline, int, T, dict[str, T]], dict[str, T]]
+
+
+def _get_text_embeddings(prompt: str, tokenizer, text_encoder, device):
+    # Tokenize text and get embeddings
+    text_inputs = tokenizer(prompt, padding='max_length', max_length=tokenizer.model_max_length, truncation=True, return_tensors='pt')
+    text_input_ids = text_inputs.input_ids
+
+    with torch.no_grad():
+        prompt_embeds = text_encoder(
+            text_input_ids.to(device),
+            output_hidden_states=True,
+        )
+
+    pooled_prompt_embeds = prompt_embeds[0]
+    prompt_embeds = prompt_embeds.hidden_states[-2]
+    if prompt == '':
+        negative_prompt_embeds = torch.zeros_like(prompt_embeds)
+        negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
+        return negative_prompt_embeds, negative_pooled_prompt_embeds
+    return prompt_embeds, pooled_prompt_embeds
+
+
+def _encode_text_sdxl(model: StableDiffusionXLPipeline, prompt: str) -> tuple[dict[str, T], T]:
+    device = model._execution_device
+    prompt_embeds, pooled_prompt_embeds, = _get_text_embeddings(prompt, model.tokenizer, model.text_encoder, device)
+    prompt_embeds_2, pooled_prompt_embeds2, = _get_text_embeddings( prompt, model.tokenizer_2, model.text_encoder_2, device)
+    prompt_embeds = torch.cat((prompt_embeds, prompt_embeds_2), dim=-1)
+    text_encoder_projection_dim = model.text_encoder_2.config.projection_dim
+    add_time_ids = model._get_add_time_ids((1024, 1024), (0, 0), (1024, 1024), torch.float16,
+                                           text_encoder_projection_dim).to(device)
+    added_cond_kwargs = {"text_embeds": pooled_prompt_embeds2, "time_ids": add_time_ids}
+    return added_cond_kwargs, prompt_embeds
+
+
+def _encode_text_sdxl_with_negative(model: StableDiffusionXLPipeline, prompt: str) -> tuple[dict[str, T], T]:
+    added_cond_kwargs, prompt_embeds = _encode_text_sdxl(model, prompt)
+    added_cond_kwargs_uncond, prompt_embeds_uncond = _encode_text_sdxl(model, "")
+    prompt_embeds = torch.cat((prompt_embeds_uncond, prompt_embeds, ))
+    added_cond_kwargs = {"text_embeds": torch.cat((added_cond_kwargs_uncond["text_embeds"], added_cond_kwargs["text_embeds"])),
+                         "time_ids": torch.cat((added_cond_kwargs_uncond["time_ids"], added_cond_kwargs["time_ids"])),}
+    return added_cond_kwargs, prompt_embeds
+
+
+def _encode_image(model: StableDiffusionXLPipeline, image: np.ndarray) -> T:
+    model.vae.to(dtype=torch.float32)
+    image = torch.from_numpy(image).float() / 255.
+    image = (image * 2 - 1).permute(2, 0, 1).unsqueeze(0)
+    latent = model.vae.encode(image.to(model.vae.device))['latent_dist'].mean * model.vae.config.scaling_factor
+    model.vae.to(dtype=torch.float16)
+    return latent
+
+
+def _next_step(model: StableDiffusionXLPipeline, model_output: T, timestep: int, sample: T) -> T:
+    timestep, next_timestep = min(timestep - model.scheduler.config.num_train_timesteps // model.scheduler.num_inference_steps, 999), timestep
+    alpha_prod_t = model.scheduler.alphas_cumprod[int(timestep)] if timestep >= 0 else model.scheduler.final_alpha_cumprod
+    alpha_prod_t_next = model.scheduler.alphas_cumprod[int(next_timestep)]
+    beta_prod_t = 1 - alpha_prod_t
+    next_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
+    next_sample_direction = (1 - alpha_prod_t_next) ** 0.5 * model_output
+    next_sample = alpha_prod_t_next ** 0.5 * next_original_sample + next_sample_direction
+    return next_sample
+
+
+def _get_noise_pred(model: StableDiffusionXLPipeline, latent: T, t: T, context: T, guidance_scale: float, added_cond_kwargs: dict[str, T]):
+    latents_input = torch.cat([latent] * 2)
+    noise_pred = model.unet(latents_input, t, encoder_hidden_states=context, added_cond_kwargs=added_cond_kwargs)["sample"]
+    noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2)
+    noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond)
+    # latents = next_step(model, noise_pred, t, latent)
+    return noise_pred
+
+
+def _ddim_loop(model: StableDiffusionXLPipeline, z0, prompt, guidance_scale) -> T:
+    all_latent = [z0]
+    added_cond_kwargs, text_embedding = _encode_text_sdxl_with_negative(model, prompt)
+    latent = z0.clone().detach().half()
+    for i in tqdm(range(model.scheduler.num_inference_steps)):
+        t = model.scheduler.timesteps[len(model.scheduler.timesteps) - i - 1]
+        noise_pred = _get_noise_pred(model, latent, t, text_embedding, guidance_scale, added_cond_kwargs)
+        latent = _next_step(model, noise_pred, t, latent)
+        all_latent.append(latent)
+    return torch.cat(all_latent).flip(0)
+
+
+def make_inversion_callback(zts, offset: int = 0) -> [T, InversionCallback]:
+
+    def callback_on_step_end(pipeline: StableDiffusionXLPipeline, i: int, t: T, callback_kwargs: dict[str, T]) -> dict[str, T]:
+        latents = callback_kwargs['latents']
+        latents[0] = zts[max(offset + 1, i + 1)].to(latents.device, latents.dtype)
+        return {'latents': latents}
+    return  zts[offset], callback_on_step_end
+
+
+@torch.no_grad()
+def ddim_inversion(model: StableDiffusionXLPipeline, x0: np.ndarray, prompt: str, num_inference_steps: int, guidance_scale,) -> T:
+    z0 = _encode_image(model, x0)
+    model.scheduler.set_timesteps(num_inference_steps, device=z0.device)
+    zs = _ddim_loop(model, z0, prompt, guidance_scale)
+    return zs

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

@@ -0,0 +1,6 @@
+diffusers
+transformers
+accelerate
+mediapy
+ipywidgets
+einops

sa_handler.py

@@ -0,0 +1,279 @@
+# 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
+    shared_score_scale: float = 1.
+    shared_score_shift: float = 0.
+    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 shifted_scaled_dot_product_attention(self, attn: attention_processor.Attention, query: T, key: T, value: T) -> T:
+        logits = torch.einsum('bhqd,bhkd->bhqk', query, key) * attn.scale
+        logits[:, :, :, query.shape[2]:] += self.shared_score_shift
+        probs = logits.softmax(-1)
+        return torch.einsum('bhqk,bhkd->bhqd', probs, value)
+
+    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.shared_score_scale)
+            value = concat_first(value, -2)
+            if self.shared_score_shift != 0:
+                hidden_states = self.shifted_scaled_dot_product_attention(attn, query, key, value,)
+            else:
+                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.shared_score_scale = style_aligned_args.shared_score_scale
+        self.shared_score_shift = style_aligned_args.shared_score_shift
+
+
+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_sd1.ipynb

@@ -0,0 +1,162 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "a885cf5d-c525-4f5b-a8e4-f67d2f699909",
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "## Copyright 2023 Google LLC"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d891d022-8979-40d4-848f-ecb84c17f12c",
+   "metadata": {
+    "jp-MarkdownHeadingCollapsed": true,
+    "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": "540d8642-c203-471c-a66d-0d43aabb0706",
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "# StyleAligned over SD1.4"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "23d54ea7-f7ab-4548-9b10-ece87216dc18",
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "from diffusers import DDIMScheduler,StableDiffusionPipeline\n",
+    "import torch\n",
+    "import mediapy\n",
+    "import sa_handler\n",
+    "import math"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "522b14e7-9768-4eaa-8433-bf88acb244c4",
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "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 = StableDiffusionPipeline.from_pretrained(\n",
+    "    \"CompVis/stable-diffusion-v1-4\",\n",
+    "    revision=\"fp16\",\n",
+    "    scheduler=scheduler\n",
+    ")\n",
+    "pipeline = pipeline.to(\"cuda\")\n",
+    "\n",
+    "handler = sa_handler.Handler(pipeline)\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",
+    "\n",
+    "handler.register(sa_args, )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5db98c81-8b72-4fc7-8cd0-65eda17198e3",
+   "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, generator=None).images\n",
+    "mediapy.show_images(images)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "afbe3876-22d9-4735-89b9-d5b5c46aea5c",
+   "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_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_transfer_sdxl.ipynb

@@ -0,0 +1,186 @@
+{
+ "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 from input image"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "483d0cf9",
+   "metadata": {},
+   "source": [
+    "#### Model Load "
+   ]
+  },
+  {
+   "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\n",
+    "import math\n",
+    "\n",
+    "\n",
+    "scheduler = DDIMScheduler(\n",
+    "    beta_start=0.00085, beta_end=0.012, beta_schedule=\"scaled_linear\",\n",
+    "    clip_sample=False, set_alpha_to_one=False)\n",
+    "\n",
+    "pipeline = StableDiffusionXLPipeline.from_pretrained(\n",
+    "    \"stabilityai/stable-diffusion-xl-base-1.0\", torch_dtype=torch.float16, variant=\"fp16\",\n",
+    "    use_safetensors=True,\n",
+    "    scheduler=scheduler\n",
+    ").to(\"cuda\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c09b1a68",
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "#### Ref image load and inversion"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f4717854",
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# DDIM inversion\n",
+    "\n",
+    "from diffusers.utils import load_image\n",
+    "import inversion\n",
+    "import numpy as np\n",
+    "\n",
+    "src_style = \"medieval painting\"\n",
+    "src_prompt = f'Man laying in a bed, {src_style}.'\n",
+    "image_path = './example_image/medieval-bed.jpeg'\n",
+    "\n",
+    "num_inference_steps = 50\n",
+    "x0 = np.array(load_image(image_path).resize((1024, 1024)))\n",
+    "zts = inversion.ddim_inversion(pipeline, x0, src_prompt, num_inference_steps, 2)\n",
+    "mediapy.show_image(x0, title=\"innput reference image\", height=256)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1751c4fe",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prompts = [\n",
+    "    src_prompt,\n",
+    "    \"A man working on a laptop\",\n",
+    "    \"A man eats pizza\",\n",
+    "    \"A woman playig on saxophone\",\n",
+    "]\n",
+    "\n",
+    "# some parameters you can adjust to control fidelity to reference\n",
+    "shared_score_shift = np.log(2)  # higher value induces higher fidelity, set 0 for no shift\n",
+    "shared_score_scale = 1.0  # higher value induces higher, set 1 for no rescale\n",
+    "\n",
+    "# for very famouse images consider supressing attention to refference, here is a configuration example:\n",
+    "# shared_score_shift = np.log(1)\n",
+    "# shared_score_scale = 0.5\n",
+    "\n",
+    "for i in range(1, len(prompts)):\n",
+    "    prompts[i] = f'{prompts[i]}, {src_style}.'\n",
+    "\n",
+    "handler = sa_handler.Handler(pipeline)\n",
+    "sa_args = sa_handler.StyleAlignedArgs(\n",
+    "    share_group_norm=True, share_layer_norm=True, share_attention=True,\n",
+    "    adain_queries=True, adain_keys=True, adain_values=False,\n",
+    "    shared_score_shift=shared_score_shift, shared_score_scale=shared_score_scale,)\n",
+    "handler.register(sa_args)\n",
+    "\n",
+    "zT, inversion_callback = inversion.make_inversion_callback(zts, offset=5)\n",
+    "\n",
+    "g_cpu = torch.Generator(device='cpu')\n",
+    "g_cpu.manual_seed(10)\n",
+    "\n",
+    "latents = torch.randn(len(prompts), 4, 128, 128, device='cpu', generator=g_cpu,\n",
+    "                      dtype=pipeline.unet.dtype,).to('cuda:0')\n",
+    "latents[0] = zT\n",
+    "\n",
+    "images_a = pipeline(prompts, latents=latents,\n",
+    "                    callback_on_step_end=inversion_callback,\n",
+    "                    num_inference_steps=num_inference_steps, guidance_scale=10.0).images\n",
+    "\n",
+    "handler.remove()\n",
+    "mediapy.show_images(images_a, titles=[p[:-(len(src_style) + 3)] for p in prompts])"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "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.10.13"
+  }
+ },
+ "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
+}