Merge branch 'dev' into flux-lora

CHANGELOG.md

@@ -1,6 +1,6 @@
 # Change Log for SD.Next
 
-## Update for 2024-09-26
+## Update for 2024-09-29
 
 - **reprocess**
   - new top-level button: reprocess your last generated image(s)  
@@ -18,6 +18,14 @@
     *note* sd/sdxl contain heavily distilled versions of reference models, so switching to reference model produces vastly different results  
   - xyz grid support for text encoder  
   - full prompt parser now correctly works with different prompts in batch  
+- [Ctrl+X](https://github.com/genforce/ctrl-x):
+  - control **structure** (*similar to controlnet*) and **appearance** (*similar to ipadapter*)  
+    without the need for extra models, all via code feed-forwards!
+  - can run in structure-only or appearance-only or both modes
+  - when providing structure and appearance input images, its best to provide a short prompts describing them  
+  - structure image can be *almost anything*: *actual photo, openpose-style stick man, 3d render, sketch, depth-map, etc.*  
+    just describe what it is in a structure prompt so it can be de-structured and correctly applied  
+  - supports sdxl in both txt2img and img2img, simply select from scripts
 - **flux**  
   - avoid unet load if unchanged  
   - mark specific unet as unavailable if load failed  
@@ -76,14 +84,16 @@
   - fix imageviewer exif parser  
   - selectable info view in image viewer, thanks @ZeldaMaster501  
 - **free-u** check if device/dtype are fft compatible and cast as necessary  
-- **experimental**  
-  - flux t5 load from gguf: requires transformers pr  
-  - rocm triton backend for flash attention, thanks @lshqqytiger  
+- **rocm**
+  - additional gpu detection and auto-config code, thanks @lshqqytiger
+  - experimental triton backend for flash attention, thanks @lshqqytiger  
 - **refactor**  
   - modularize main process loop  
   - massive log cleanup  
   - full lint pass  
   - improve inference mode handling  
+- **experimental**  
+  - flux t5 load from gguf: requires transformers pr  
 
 
 ## Update for 2024-09-13
@@ -191,7 +201,9 @@ Examples:
 - **prompt enhance**: improve quality and/or verbosity of your prompts  
   simply select in *scripts -> prompt enhance*
   uses [gokaygokay/Flux-Prompt-Enhance](https://huggingface.co/gokaygokay/Flux-Prompt-Enhance) model  
-- **decode** auto-set upcast if first decode fails  
+- **decode**
+  - auto-set upcast if first decode fails  
+  - restore dtype on upcast  
 - **taesd** configurable number of layers  
   can be used to speed-up taesd decoding by reducing number of ops  
   e.g. if generating 1024px image, reducing layers by 1 will result in preview being 512px  

installer.py

@@ -112,7 +112,7 @@ def setup_logging():
     }))
     logging.basicConfig(level=logging.ERROR, format='%(asctime)s | %(name)s | %(levelname)s | %(module)s | %(message)s', handlers=[logging.NullHandler()]) # redirect default logger to null
     pretty_install(console=console)
-    traceback_install(console=console, extra_lines=1, max_frames=10, width=console.width, word_wrap=False, indent_guides=False, suppress=[])
+    traceback_install(console=console, extra_lines=1, max_frames=16, width=console.width, word_wrap=False, indent_guides=False, suppress=[])
     while log.hasHandlers() and len(log.handlers) > 0:
         log.removeHandler(log.handlers[0])
 
@@ -477,27 +477,21 @@ def install_rocm_zluda():
     try:
         amd_gpus = rocm.get_agents()
         if len(amd_gpus) == 0:
-            if sys.platform == "win32":
-                log.warning('You do not have perl or any AMDGPUs. The installer may select a wrong device as compute device.')
-                log.info('ROCm: no agent was found')
-            else:
-                log.warning('ROCm: no agent was found')
+            log.warning('ROCm: no agent was found')
         else:
             log.info(f'ROCm: agents={[gpu.name for gpu in amd_gpus]}')
             if args.device_id is None:
-                device = amd_gpus[0]
+                index = 0
                 for idx, gpu in enumerate(amd_gpus):
-                    if gpu.arch == rocm.MicroArchitecture.RDNA:
-                        device = gpu
-                        os.environ.setdefault('HIP_VISIBLE_DEVICES', str(idx))
-                        # if os.environ.get('TENSORFLOW_PACKAGE') == 'tensorflow-rocm': # do not use tensorflow-rocm for navi 3x
-                        #    os.environ['TENSORFLOW_PACKAGE'] = 'tensorflow==2.13.0'
-                        if not device.is_apu:
-                            # although apu was found, there can be a dedicated card. do not break loop.
-                            # if no dedicated card was found, apu will be used.
-                            break
-                    else:
-                        log.debug(f'ROCm: HSA_OVERRIDE_GFX_VERSION auto config skipped for {gpu.name}')
+                    index = idx
+                    # if gpu.name.startswith('gfx11') and os.environ.get('TENSORFLOW_PACKAGE') == 'tensorflow-rocm': # do not use tensorflow-rocm for navi 3x
+                    #    os.environ['TENSORFLOW_PACKAGE'] = 'tensorflow==2.13.0'
+                    if not gpu.is_apu:
+                        # although apu was found, there can be a dedicated card. do not break loop.
+                        # if no dedicated card was found, apu will be used.
+                        break
+                os.environ.setdefault('HIP_VISIBLE_DEVICES', str(index))
+                device = amd_gpus[index]
             else:
                 device_id = int(args.device_id)
                 if device_id < len(amd_gpus):
@@ -571,7 +565,9 @@ def install_rocm_zluda():
             log.debug(f'ROCm hipBLASLt: arch={device.name} available={device.blaslt_supported}')
             rocm.set_blaslt_enabled(device.blaslt_supported)
 
-    if device is not None:
+    if device is None:
+        log.debug('ROCm: HSA_OVERRIDE_GFX_VERSION auto config skipped')
+    else:
         os.environ.setdefault('HSA_OVERRIDE_GFX_VERSION', device.get_gfx_version())
 
     return torch_command

modules/ctrlx/__init__.py

@@ -0,0 +1,662 @@
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from diffusers import StableDiffusionXLPipeline
+from diffusers.image_processor import PipelineImageInput
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img import rescale_noise_cfg, retrieve_latents, retrieve_timesteps
+from diffusers.utils import BaseOutput, deprecate
+from diffusers.utils.torch_utils import randn_tensor
+import numpy as np
+import PIL
+import torch
+from .sdxl import register_attr
+from .media import preprocess
+from .utils import batch_dict_to_tensor, batch_tensor_to_dict, noise_prev, noise_t2t
+
+
+BATCH_ORDER = [
+    "structure_uncond", "appearance_uncond", "uncond", "structure_cond", "appearance_cond", "cond",
+]
+
+
+def get_last_control_i(control_schedule, num_inference_steps):
+    if control_schedule is None:
+        return num_inference_steps, num_inference_steps
+
+    def max_(l):
+        if len(l) == 0:
+            return 0.0
+        return max(l)
+
+    structure_max = 0.0
+    appearance_max = 0.0
+    for block in control_schedule.values():
+        if isinstance(block, list):  # Handling mid_block
+            block = {0: block}
+        for layer in block.values():
+            structure_max = max(structure_max, max_(layer[0] + layer[1]))
+            appearance_max = max(appearance_max, max_(layer[2]))
+
+    structure_i = round(num_inference_steps * structure_max)
+    appearance_i = round(num_inference_steps * appearance_max)
+    return structure_i, appearance_i
+
+
+@dataclass
+class CtrlXStableDiffusionXLPipelineOutput(BaseOutput):
+    images: Union[List[PIL.Image.Image], np.ndarray] = None
+    structures: Union[List[PIL.Image.Image], np.ndarray] = None
+    appearances: Union[List[PIL.Image.Image], np.ndarray] = None
+
+
+class CtrlXStableDiffusionXLPipeline(StableDiffusionXLPipeline):  # diffusers==0.28.0
+
+    def prepare_latents(
+        self, image, batch_size, num_images_per_prompt, num_channels_latents, height, width,
+        dtype, device, generator=None, noise=None,
+    ):
+        batch_size = batch_size * num_images_per_prompt
+
+        if noise is None:
+            shape = (
+                batch_size,
+                num_channels_latents,
+                height // self.vae_scale_factor,
+                width // self.vae_scale_factor
+            )
+            noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+            noise = noise * self.scheduler.init_noise_sigma  # Starting noise, need to scale
+        else:
+            noise = noise.to(device)
+
+        if image is None:
+            return noise, None
+
+        if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
+            raise ValueError(
+                f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
+            )
+
+        # Offload text encoder if `enable_model_cpu_offload` was enabled
+        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+            self.text_encoder_2.to("cpu")
+            torch.cuda.empty_cache()
+
+        image = image.to(device=device, dtype=dtype)
+
+        if image.shape[1] == 4:  # Image already in latents form
+            init_latents = image
+
+        else:
+            # Make sure the VAE is in float32 mode, as it overflows in float16
+            if self.vae.config.force_upcast:
+                image = image.to(torch.float32)
+                self.vae.to(torch.float32)
+
+            if isinstance(generator, list) and len(generator) != batch_size:
+                raise ValueError(
+                    f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                    f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+                )
+            elif isinstance(generator, list):
+                init_latents = [
+                    retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i])
+                    for i in range(batch_size)
+                ]
+                init_latents = torch.cat(init_latents, dim=0)
+            else:
+                init_latents = retrieve_latents(self.vae.encode(image), generator=generator)
+
+            if self.vae.config.force_upcast:
+                self.vae.to(dtype)
+
+            init_latents = init_latents.to(dtype)
+            init_latents = self.vae.config.scaling_factor * init_latents
+
+        if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0:
+            # Expand init_latents for batch_size
+            additional_image_per_prompt = batch_size // init_latents.shape[0]
+            init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0)
+        elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0:
+            raise ValueError(
+                f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."
+            )
+        else:
+            init_latents = torch.cat([init_latents], dim=0)
+
+        return noise, init_latents
+
+    @property
+    def structure_guidance_scale(self):
+        return self._guidance_scale if self._structure_guidance_scale is None else self._structure_guidance_scale
+
+    @property
+    def appearance_guidance_scale(self):
+        return self._guidance_scale if self._appearance_guidance_scale is None else self._appearance_guidance_scale
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,  # TODO: Support prompt_2 and negative_prompt_2
+        structure_prompt: Optional[Union[str, List[str]]] = None,
+        appearance_prompt: Optional[Union[str, List[str]]] = None,
+        structure_image: Optional[PipelineImageInput] = None,
+        appearance_image: Optional[PipelineImageInput] = None,
+        num_inference_steps: int = 50,
+        timesteps: List[int] = None,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        positive_prompt: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        guidance_scale: float = 5.0,
+        structure_guidance_scale: Optional[float] = None,
+        appearance_guidance_scale: Optional[float] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        structure_latents: Optional[torch.Tensor] = None,
+        appearance_latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,  # Positive prompt is concatenated with prompt, so no embeddings
+        structure_prompt_embeds: Optional[torch.Tensor] = None,
+        appearance_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        pooled_prompt_embeds: Optional[torch.Tensor] = None,
+        structure_pooled_prompt_embeds: Optional[torch.Tensor] = None,
+        appearance_pooled_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
+        control_schedule: Optional[Dict] = None,
+        self_recurrence_schedule: Optional[List[int]] = [],  # Format: [(start, end, num_repeat)]
+        decode_structure: Optional[bool] = True,
+        decode_appearance: Optional[bool] = True,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        guidance_rescale: float = 0.0,
+        original_size: Tuple[int, int] = None,
+        crops_coords_top_left: Tuple[int, int] = (0, 0),
+        target_size: Tuple[int, int] = None,
+        clip_skip: Optional[int] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        **kwargs,
+    ):
+        # TODO: Add function argument documentation
+
+        callback = kwargs.pop("callback", None)
+        callback_steps = kwargs.pop("callback_steps", None)
+
+        if callback is not None:
+            deprecate(
+                "callback",
+                "1.0.0",
+                "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
+            )
+        if callback_steps is not None:
+            deprecate(
+                "callback_steps",
+                "1.0.0",
+                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
+            )
+
+        # 0. Default height and width to U-Net
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+        original_size = original_size or (height, width)
+        target_size = target_size or (height, width)
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(  # TODO: Custom check_inputs for our method
+            prompt,
+            None,  # prompt_2
+            height,
+            width,
+            callback_steps,
+            negative_prompt = negative_prompt,
+            negative_prompt_2 = None,  # 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,
+            callback_on_step_end_tensor_inputs = callback_on_step_end_tensor_inputs,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._structure_guidance_scale = structure_guidance_scale
+        self._appearance_guidance_scale = appearance_guidance_scale
+        self._guidance_rescale = guidance_rescale
+        self._clip_skip = clip_skip
+        self._cross_attention_kwargs = cross_attention_kwargs
+        self._denoising_end = None  # denoising_end
+        self._denoising_start = None  # denoising_start
+        self._interrupt = False
+
+        # 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]
+
+        if batch_size * num_images_per_prompt != 1:
+            raise ValueError(
+                f"Pipeline currently does not support batch_size={batch_size} and num_images_per_prompt=1. "
+                "Effective batch size (batch_size * num_images_per_prompt) must be 1."
+            )
+
+        device = self._execution_device
+
+        # 3. Encode input prompt
+        text_encoder_lora_scale = (
+            self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
+        )
+
+        if positive_prompt is not None and positive_prompt != "":
+            prompt = prompt + ", " + positive_prompt  # Add positive prompt with comma
+            # By default, only add positive prompt to the appearance prompt and not the structure prompt
+            if appearance_prompt is not None and appearance_prompt != "":
+                appearance_prompt = appearance_prompt + ", " + positive_prompt
+
+        (
+            prompt_embeds_,
+            negative_prompt_embeds,
+            pooled_prompt_embeds_,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt = prompt,
+            prompt_2 = None,  # prompt_2
+            device = device,
+            num_images_per_prompt = num_images_per_prompt,
+            do_classifier_free_guidance = True,  # self.do_classifier_free_guidance, TODO: Support no CFG
+            negative_prompt = negative_prompt,
+            negative_prompt_2 = None,  # 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 = self.clip_skip,
+        )
+        prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds_], dim=0).to(device)
+        add_text_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds_], dim=0).to(device)
+
+        # 3.1. Structure prompt embeddings
+        if structure_prompt is not None and structure_prompt != "":
+            (
+                structure_prompt_embeds,
+                negative_structure_prompt_embeds,
+                structure_pooled_prompt_embeds,
+                negative_structure_pooled_prompt_embeds,
+            ) = self.encode_prompt(
+                prompt = structure_prompt,
+                prompt_2 = None,  # prompt_2
+                device = device,
+                num_images_per_prompt = num_images_per_prompt,
+                do_classifier_free_guidance = True,  # self.do_classifier_free_guidance, TODO: Support no CFG
+                negative_prompt = negative_prompt if structure_image is None else "",
+                negative_prompt_2 = None,  # negative_prompt_2
+                prompt_embeds = structure_prompt_embeds,
+                negative_prompt_embeds = None,  # negative_prompt_embeds
+                pooled_prompt_embeds = structure_pooled_prompt_embeds,
+                negative_pooled_prompt_embeds = None,  # negative_pooled_prompt_embeds
+                lora_scale = text_encoder_lora_scale,
+                clip_skip = self.clip_skip,
+            )
+            structure_prompt_embeds = torch.cat(
+                [negative_structure_prompt_embeds, structure_prompt_embeds], dim=0
+            ).to(device)
+            structure_add_text_embeds = torch.cat(
+                [negative_structure_pooled_prompt_embeds, structure_pooled_prompt_embeds], dim=0
+            ).to(device)
+        else:
+            structure_prompt_embeds = prompt_embeds
+            structure_add_text_embeds = add_text_embeds
+
+        # 3.2. Appearance prompt embeddings
+        if appearance_prompt is not None and appearance_prompt != "":
+            (
+                appearance_prompt_embeds,
+                negative_appearance_prompt_embeds,
+                appearance_pooled_prompt_embeds,
+                negative_appearance_pooled_prompt_embeds,
+            ) = self.encode_prompt(
+                prompt = appearance_prompt,
+                prompt_2 = None,  # prompt_2
+                device = device,
+                num_images_per_prompt = num_images_per_prompt,
+                do_classifier_free_guidance = True,  # self.do_classifier_free_guidance, TODO: Support no CFG
+                negative_prompt = negative_prompt if appearance_image is None else "",
+                negative_prompt_2 = None,  # negative_prompt_2
+                prompt_embeds = appearance_prompt_embeds,
+                negative_prompt_embeds = None,  # negative_prompt_embeds
+                pooled_prompt_embeds = appearance_pooled_prompt_embeds,  # pooled_prompt_embeds
+                negative_pooled_prompt_embeds = None,  # negative_pooled_prompt_embeds
+                lora_scale = text_encoder_lora_scale,
+                clip_skip = self.clip_skip,
+            )
+            appearance_prompt_embeds = torch.cat(
+                [negative_appearance_prompt_embeds, appearance_prompt_embeds], dim=0
+            ).to(device)
+            appearance_add_text_embeds = torch.cat(
+                [negative_appearance_pooled_prompt_embeds, appearance_pooled_prompt_embeds], dim=0
+            ).to(device)
+        else:
+            appearance_prompt_embeds = prompt_embeds
+            appearance_add_text_embeds = add_text_embeds
+
+        # 3.3. Prepare added time ids & embeddings, TODO: Support no CFG
+        if self.text_encoder_2 is None:
+            text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
+        else:
+            text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
+
+        add_time_ids = self._get_add_time_ids(
+            original_size,
+            crops_coords_top_left,
+            target_size,
+            dtype = prompt_embeds.dtype,
+            text_encoder_projection_dim = text_encoder_projection_dim,
+        )
+        negative_add_time_ids = add_time_ids
+        add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0).to(device)
+
+        # 4. Prepare timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+
+        latents, _ = self.prepare_latents(
+            None, batch_size, num_images_per_prompt, num_channels_latents, height, width,
+            prompt_embeds.dtype, device, generator, latents
+        )
+
+        if structure_image is not None:
+            structure_image = preprocess(  # Center crop + resize
+                structure_image, self.image_processor, height=height, width=width, resize_mode="crop"
+            )
+            _, clean_structure_latents = self.prepare_latents(
+                structure_image, batch_size, num_images_per_prompt, num_channels_latents, height, width,
+                prompt_embeds.dtype, device, generator, structure_latents,
+            )
+        else:
+            clean_structure_latents = None
+        structure_latents = latents if structure_latents is None else structure_latents
+
+        if appearance_image is not None:
+            appearance_image = preprocess(  # Center crop + resize
+                appearance_image, self.image_processor, height=height, width=width, resize_mode="crop"
+            )
+            _, clean_appearance_latents = self.prepare_latents(
+                appearance_image, batch_size, num_images_per_prompt, num_channels_latents, height, width,
+                prompt_embeds.dtype, device, generator, appearance_latents,
+            )
+        else:
+            clean_appearance_latents = None
+        appearance_latents = latents if appearance_latents is None else appearance_latents
+
+        # 6. Prepare extra step kwargs
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Denoising loop
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+
+        # 7.1 Apply denoising_end
+        def denoising_value_valid(dnv):
+            return isinstance(self.denoising_end, float) and 0 < dnv < 1
+
+        if (
+            self.denoising_end is not None
+            and self.denoising_start is not None
+            and denoising_value_valid(self.denoising_end)
+            and denoising_value_valid(self.denoising_start)
+            and self.denoising_start >= self.denoising_end
+        ):
+            raise ValueError(
+                f"`denoising_start`: {self.denoising_start} cannot be larger than or equal to `denoising_end`: "
+                + f" {self.denoising_end} when using type float."
+            )
+        elif self.denoising_end is not None and denoising_value_valid(self.denoising_end):
+            discrete_timestep_cutoff = int(
+                round(
+                    self.scheduler.config.num_train_timesteps
+                    - (self.denoising_end * self.scheduler.config.num_train_timesteps)
+                )
+            )
+            num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))
+            timesteps = timesteps[:num_inference_steps]
+
+        # 7.2 Optionally get guidance scale embedding
+        timestep_cond = None
+        if self.unet.config.time_cond_proj_dim is not None:  # TODO: Make guidance scale embedding work with batch_order
+            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt)
+            timestep_cond = self.get_guidance_scale_embedding(
+                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
+            ).to(device=device, dtype=latents.dtype)
+
+        # 7.3 Get batch order
+        batch_order = deepcopy(BATCH_ORDER)
+        if structure_image is not None:  # If image is provided, not generating, so no CFG needed
+            batch_order.remove("structure_uncond")
+        if appearance_image is not None:
+            batch_order.remove("appearance_uncond")
+
+        structure_control_stop_i, appearance_control_stop_i = get_last_control_i(control_schedule, num_inference_steps)
+        if self_recurrence_schedule is None or len(self_recurrence_schedule) == 0:
+            self_recurrence_schedule = [0] * num_inference_steps
+
+        self._num_timesteps = len(timesteps)
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                if i == structure_control_stop_i:  # If not generating structure/appearance, drop after last control
+                    if "structure_uncond" not in batch_order:
+                        batch_order.remove("structure_cond")
+                if i == appearance_control_stop_i:
+                    if "appearance_uncond" not in batch_order:
+                        batch_order.remove("appearance_cond")
+
+                register_attr(self, t=t.item(), do_control=True, batch_order=batch_order)
+
+                # TODO: For now, assume we are doing classifier-free guidance, support no CF-guidance later
+                latent_model_input = self.scheduler.scale_model_input(latents, t)
+                structure_latent_model_input = self.scheduler.scale_model_input(structure_latents, t)
+                appearance_latent_model_input = self.scheduler.scale_model_input(appearance_latents, t)
+
+                all_latent_model_input = {
+                    "structure_uncond": structure_latent_model_input[0:1],
+                    "appearance_uncond": appearance_latent_model_input[0:1],
+                    "uncond": latent_model_input[0:1],
+                    "structure_cond": structure_latent_model_input[0:1],
+                    "appearance_cond": appearance_latent_model_input[0:1],
+                    "cond": latent_model_input[0:1],
+                }
+                all_prompt_embeds = {
+                    "structure_uncond": structure_prompt_embeds[0:1],
+                    "appearance_uncond": appearance_prompt_embeds[0:1],
+                    "uncond": prompt_embeds[0:1],
+                    "structure_cond": structure_prompt_embeds[1:2],
+                    "appearance_cond": appearance_prompt_embeds[1:2],
+                    "cond": prompt_embeds[1:2],
+                }
+                all_add_text_embeds = {
+                    "structure_uncond": structure_add_text_embeds[0:1],
+                    "appearance_uncond": appearance_add_text_embeds[0:1],
+                    "uncond": add_text_embeds[0:1],
+                    "structure_cond": structure_add_text_embeds[1:2],
+                    "appearance_cond": appearance_add_text_embeds[1:2],
+                    "cond": add_text_embeds[1:2],
+                }
+                all_time_ids = {
+                    "structure_uncond": add_time_ids[0:1],
+                    "appearance_uncond": add_time_ids[0:1],
+                    "uncond": add_time_ids[0:1],
+                    "structure_cond": add_time_ids[1:2],
+                    "appearance_cond": add_time_ids[1:2],
+                    "cond": add_time_ids[1:2],
+                }
+
+                concat_latent_model_input = batch_dict_to_tensor(all_latent_model_input, batch_order)
+                concat_prompt_embeds = batch_dict_to_tensor(all_prompt_embeds, batch_order)
+                concat_add_text_embeds = batch_dict_to_tensor(all_add_text_embeds, batch_order)
+                concat_add_time_ids = batch_dict_to_tensor(all_time_ids, batch_order)
+
+                # Predict the noise residual
+                added_cond_kwargs = {"text_embeds": concat_add_text_embeds, "time_ids": concat_add_time_ids}
+
+                concat_noise_pred = self.unet(
+                    concat_latent_model_input,
+                    t,
+                    encoder_hidden_states = concat_prompt_embeds,
+                    timestep_cond = timestep_cond,
+                    cross_attention_kwargs = self.cross_attention_kwargs,
+                    added_cond_kwargs = added_cond_kwargs,
+                ).sample
+                all_noise_pred = batch_tensor_to_dict(concat_noise_pred, batch_order)
+
+                # Classifier-free guidance, TODO: Support no CFG
+                noise_pred = all_noise_pred["uncond"] +\
+                    self.guidance_scale * (all_noise_pred["cond"] - all_noise_pred["uncond"])
+
+                structure_noise_pred = all_noise_pred["structure_cond"]\
+                    if "structure_cond" in batch_order else noise_pred
+                if "structure_uncond" in all_noise_pred:
+                    structure_noise_pred = all_noise_pred["structure_uncond"] +\
+                        self.structure_guidance_scale * (structure_noise_pred - all_noise_pred["structure_uncond"])
+
+                appearance_noise_pred = all_noise_pred["appearance_cond"]\
+                    if "appearance_cond" in batch_order else noise_pred
+                if "appearance_uncond" in all_noise_pred:
+                    appearance_noise_pred = all_noise_pred["appearance_uncond"] +\
+                        self.appearance_guidance_scale * (appearance_noise_pred - all_noise_pred["appearance_uncond"])
+
+                if self.guidance_rescale > 0.0:
+                    noise_pred = rescale_noise_cfg(
+                        noise_pred, all_noise_pred["cond"], guidance_rescale=self.guidance_rescale
+                    )
+                    if "structure_uncond" in all_noise_pred:
+                        structure_noise_pred = rescale_noise_cfg(
+                            structure_noise_pred, all_noise_pred["structure_cond"],
+                            guidance_rescale=self.guidance_rescale
+                        )
+                    if "appearance_uncond" in all_noise_pred:
+                        appearance_noise_pred = rescale_noise_cfg(
+                            appearance_noise_pred, all_noise_pred["appearance_cond"],
+                            guidance_rescale=self.guidance_rescale
+                        )
+
+                # Compute the previous noisy sample x_t -> x_t-1
+                concat_noise_pred = torch.cat(
+                    [structure_noise_pred, appearance_noise_pred, noise_pred], dim=0,
+                )
+                concat_latents = torch.cat(
+                    [structure_latents, appearance_latents, latents], dim=0,
+                )
+                structure_latents, appearance_latents, latents = self.scheduler.step(
+                    concat_noise_pred, t, concat_latents, **extra_step_kwargs,
+                ).prev_sample.chunk(3)
+
+                if clean_structure_latents is not None:
+                    structure_latents = noise_prev(self.scheduler, t, clean_structure_latents)
+                if clean_appearance_latents is not None:
+                    appearance_latents = noise_prev(self.scheduler, t, clean_appearance_latents)
+
+                # Self-recurrence
+                for _ in range(self_recurrence_schedule[i]):
+                    if hasattr(self.scheduler, "_step_index"):  # For fancier schedulers
+                        self.scheduler._step_index -= 1  # TODO: Does this actually work?
+
+                    t_prev = 0 if i + 1 >= num_inference_steps else timesteps[i + 1]
+                    latents = noise_t2t(self.scheduler, t_prev, t, latents)
+                    latent_model_input = torch.cat([latents] * 2)
+
+                    register_attr(self, t=t.item(), do_control=False, batch_order=["uncond", "cond"])
+
+                    # Predict the noise residual
+                    added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
+                    noise_pred_uncond, noise_pred_ = self.unet(
+                        latent_model_input,
+                        t,
+                        encoder_hidden_states = prompt_embeds,
+                        timestep_cond = timestep_cond,
+                        cross_attention_kwargs = self.cross_attention_kwargs,
+                        added_cond_kwargs = added_cond_kwargs,
+                    ).sample.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_ - noise_pred_uncond)
+
+                    if self.guidance_rescale > 0.0:
+                        noise_pred = rescale_noise_cfg(noise_pred, noise_pred_, guidance_rescale=self.guidance_rescale)
+
+                    latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+
+                # Callbacks
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+                    add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds)
+                    negative_pooled_prompt_embeds = callback_outputs.pop("negative_pooled_prompt_embeds", negative_pooled_prompt_embeds)
+                    add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids)
+                    # add_neg_time_ids = callback_outputs.pop("add_neg_time_ids", add_neg_time_ids)
+
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        callback(step_idx, t, latents)
+
+        # "Reconstruction"
+        if clean_structure_latents is not None:
+            structure_latents = clean_structure_latents
+        if clean_appearance_latents is not None:
+            appearance_latents = clean_appearance_latents
+
+        # For passing important information onto the refiner
+        self.refiner_args = {"latents": latents.detach(), "prompt": prompt, "negative_prompt": negative_prompt}
+
+        if output_type != "latent":
+            # Make sure the VAE is in float32 mode, as it overflows in float16
+            if self.vae.config.force_upcast:
+                self.upcast_vae()
+                vae_dtype = next(iter(self.vae.post_quant_conv.parameters())).dtype
+                latents = latents.to(vae_dtype)
+                structure_latents = structure_latents.to(vae_dtype)
+                appearance_latents = appearance_latents.to(vae_dtype)
+
+            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+            if decode_structure:
+                structure = self.vae.decode(structure_latents / self.vae.config.scaling_factor, return_dict=False)[0]
+                structure = self.image_processor.postprocess(structure, output_type=output_type)
+            else:
+                structure = structure_latents
+            if decode_appearance:
+                appearance = self.vae.decode(appearance_latents / self.vae.config.scaling_factor, return_dict=False)[0]
+                appearance = self.image_processor.postprocess(appearance, output_type=output_type)
+            else:
+                appearance = appearance_latents
+
+            # Cast back to fp16 if needed
+            if self.vae.config.force_upcast:
+                self.vae.to(dtype=torch.float16)
+
+        else:
+            # combined = torch.cat([latents, structure_latents, appearance_latents], dim=0)
+            # return CtrlXStableDiffusionXLPipelineOutput(images=combined)
+            return CtrlXStableDiffusionXLPipelineOutput(images=latents, structures=structure_latents, appearances=appearance_latents)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image, structure, appearance)
+
+        return CtrlXStableDiffusionXLPipelineOutput(images=image, structures=structure, appearances=appearance)

modules/ctrlx/features.py

@@ -0,0 +1,70 @@
+import torch.nn.functional as F
+from .utils import batch_dict_to_tensor, batch_tensor_to_dict
+
+
+def get_schedule(timesteps, schedule):
+    end = round(len(timesteps) * schedule)
+    timesteps = timesteps[:end]
+    return timesteps
+
+
+def get_elem(l, i, default=0.0):
+    if i >= len(l):
+        return default
+    return l[i]
+
+
+def pad_list(l_1, l_2, pad=0.0):
+    max_len = max(len(l_1), len(l_2))
+    l_1 = l_1 + [pad] * (max_len - len(l_1))
+    l_2 = l_2 + [pad] * (max_len - len(l_2))
+    return l_1, l_2
+
+
+def normalize(x, dim):
+    x_mean = x.mean(dim=dim, keepdim=True)
+    x_std = x.std(dim=dim, keepdim=True)
+    x_normalized = (x - x_mean) / x_std
+    return x_normalized
+
+
+# https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
+def appearance_mean_std(q_c_normed, k_s_normed, v_s):  # c: content, s: style
+    q_c = q_c_normed  # q_c and k_s must be projected from normalized features
+    k_s = k_s_normed
+    mean = F.scaled_dot_product_attention(q_c, k_s, v_s)  # Use scaled_dot_product_attention for efficiency
+    std = (F.scaled_dot_product_attention(q_c, k_s, v_s.square()) - mean.square()).relu().sqrt()
+
+    return mean, std
+
+
+def feature_injection(features, batch_order):
+    assert features.shape[0] % len(batch_order) == 0
+    features_dict = batch_tensor_to_dict(features, batch_order)
+    features_dict["cond"] = features_dict["structure_cond"]
+    features = batch_dict_to_tensor(features_dict, batch_order)
+    return features
+
+
+def appearance_transfer(features, q_normed, k_normed, batch_order, v=None, reshape_fn=None):
+    assert features.shape[0] % len(batch_order) == 0
+
+    features_dict = batch_tensor_to_dict(features, batch_order)
+    q_normed_dict = batch_tensor_to_dict(q_normed, batch_order)
+    k_normed_dict = batch_tensor_to_dict(k_normed, batch_order)
+    v_dict = features_dict
+    if v is not None:
+        v_dict = batch_tensor_to_dict(v, batch_order)
+
+    mean_cond, std_cond = appearance_mean_std(
+        q_normed_dict["cond"], k_normed_dict["appearance_cond"], v_dict["appearance_cond"],
+    )
+
+    if reshape_fn is not None:
+        mean_cond = reshape_fn(mean_cond)
+        std_cond = reshape_fn(std_cond)
+
+    features_dict["cond"] = std_cond * normalize(features_dict["cond"], dim=-2) + mean_cond
+
+    features = batch_dict_to_tensor(features_dict, batch_order)
+    return features

modules/ctrlx/media.py

@@ -0,0 +1,21 @@
+import numpy as np
+import torch
+import torchvision.transforms.functional as vF
+import PIL
+
+
+JPEG_QUALITY = 95
+
+
+def preprocess(image, processor, **kwargs):
+    if isinstance(image, PIL.Image.Image):
+        pass
+    elif isinstance(image, np.ndarray):
+        image = PIL.Image.fromarray(image)
+    elif isinstance(image, torch.Tensor):
+        image = vF.to_pil_image(image)
+    else:
+        raise TypeError(f"Image must be of type PIL.Image, np.ndarray, or torch.Tensor, got {type(image)} instead.")
+
+    image = processor.preprocess(image, **kwargs)
+    return image

modules/ctrlx/sdxl.py

@@ -0,0 +1,299 @@
+from types import MethodType
+from typing import Optional
+from diffusers.models.attention_processor import Attention
+import torch
+import torch.nn.functional as F
+from .features import feature_injection, normalize, appearance_transfer, get_elem, get_schedule
+
+
+def get_control_config(structure_schedule, appearance_schedule):
+    s = structure_schedule
+    a = appearance_schedule
+
+    control_config =\
+f"""control_schedule:
+    #       structure_conv   structure_attn   appearance_attn  conv/attn
+    encoder:                                                # (num layers)
+        0: [[             ], [             ], [             ]]  # 2/0
+        1: [[             ], [             ], [{a}, {a}     ]]  # 2/2
+        2: [[             ], [             ], [{a}, {a}     ]]  # 2/2
+    middle: [[            ], [             ], [             ]]  # 2/1
+    decoder:
+        0: [[{s}          ], [{s}, {s}, {s}], [0.0, {a}, {a}]]  # 3/3
+        1: [[             ], [             ], [{a}, {a}     ]]  # 3/3
+        2: [[             ], [             ], [             ]]  # 3/0
+
+control_target:
+    - [output_tensor]  # structure_conv   choices: {{hidden_states, output_tensor}}
+    - [query, key]     # structure_attn   choices: {{query, key, value}}
+    - [before]         # appearance_attn  choices: {{before, value, after}}
+
+self_recurrence_schedule:
+    - [0.1, 0.5, 2]  # format: [start, end, num_recurrence]"""
+
+    return control_config
+
+
+def convolution_forward(  # From <class 'diffusers.models.resnet.ResnetBlock2D'>, forward (diffusers==0.28.0)
+    self,
+    input_tensor: torch.Tensor,
+    temb: torch.Tensor,
+    *args,
+    **kwargs,
+) -> torch.Tensor:
+    do_structure_control = self.do_control and self.t in self.structure_schedule
+
+    hidden_states = input_tensor
+
+    hidden_states = self.norm1(hidden_states)
+    hidden_states = self.nonlinearity(hidden_states)
+
+    if self.upsample is not None:
+        # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
+        if hidden_states.shape[0] >= 64:
+            input_tensor = input_tensor.contiguous()
+            hidden_states = hidden_states.contiguous()
+        input_tensor = self.upsample(input_tensor)
+        hidden_states = self.upsample(hidden_states)
+    elif self.downsample is not None:
+        input_tensor = self.downsample(input_tensor)
+        hidden_states = self.downsample(hidden_states)
+
+    hidden_states = self.conv1(hidden_states)
+
+    if self.time_emb_proj is not None:
+        if not self.skip_time_act:
+            temb = self.nonlinearity(temb)
+        temb = self.time_emb_proj(temb)[:, :, None, None]
+
+    if self.time_embedding_norm == "default":
+        if temb is not None:
+            hidden_states = hidden_states + temb
+        hidden_states = self.norm2(hidden_states)
+    elif self.time_embedding_norm == "scale_shift":
+        if temb is None:
+            raise ValueError(
+                f" `temb` should not be None when `time_embedding_norm` is {self.time_embedding_norm}"
+            )
+        time_scale, time_shift = torch.chunk(temb, 2, dim=1)
+        hidden_states = self.norm2(hidden_states)
+        hidden_states = hidden_states * (1 + time_scale) + time_shift
+    else:
+        hidden_states = self.norm2(hidden_states)
+
+    hidden_states = self.nonlinearity(hidden_states)
+
+    hidden_states = self.dropout(hidden_states)
+    hidden_states = self.conv2(hidden_states)
+
+    # Feature injection and AdaIN (hidden_states)
+    if do_structure_control and "hidden_states" in self.structure_target:
+        hidden_states = feature_injection(hidden_states, batch_order=self.batch_order)
+
+    if self.conv_shortcut is not None:
+        input_tensor = self.conv_shortcut(input_tensor)
+
+    output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
+
+    # Feature injection and AdaIN (output_tensor)
+    if do_structure_control and "output_tensor" in self.structure_target:
+        output_tensor = feature_injection(output_tensor, batch_order=self.batch_order)
+
+    return output_tensor
+
+
+class AttnProcessor2_0:  # From <class 'diffusers.models.attention_processor.AttnProcessor2_0'> (diffusers==0.28.0)
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        temb: Optional[torch.FloatTensor] = None,
+        *args,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        do_structure_control = attn.do_control and attn.t in attn.structure_schedule
+        do_appearance_control = attn.do_control and attn.t in attn.appearance_schedule
+
+        residual = hidden_states
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = 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)
+
+        no_encoder_hidden_states = encoder_hidden_states is None
+        if no_encoder_hidden_states:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        if do_appearance_control:  # Assume we only have this for self attention
+            hidden_states_normed = normalize(hidden_states, dim=-2)  # B H D C
+            encoder_hidden_states_normed = normalize(encoder_hidden_states, dim=-2)
+
+            query_normed = attn.to_q(hidden_states_normed)
+            key_normed = attn.to_k(encoder_hidden_states_normed)
+
+            inner_dim = key_normed.shape[-1]
+            head_dim = inner_dim // attn.heads
+            query_normed = query_normed.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            key_normed = key_normed.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+            # Match query and key injection with structure injection (if injection is happening this layer)
+            if do_structure_control:
+                if "query" in attn.structure_target:
+                    query_normed = feature_injection(query_normed, batch_order=attn.batch_order)
+                if "key" in attn.structure_target:
+                    key_normed = feature_injection(key_normed, batch_order=attn.batch_order)
+
+        # Appearance transfer (before)
+        if do_appearance_control and "before" in attn.appearance_target:
+            hidden_states = hidden_states.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            hidden_states = appearance_transfer(hidden_states, query_normed, key_normed, batch_order=attn.batch_order)
+            hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+
+            if no_encoder_hidden_states:
+                encoder_hidden_states = hidden_states
+            elif attn.norm_cross:
+                encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        query = attn.to_q(hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_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)
+
+        # Feature injection (query, key, and/or value)
+        if do_structure_control:
+            if "query" in attn.structure_target:
+                query = feature_injection(query, batch_order=attn.batch_order)
+            if "key" in attn.structure_target:
+                key = feature_injection(key, batch_order=attn.batch_order)
+            if "value" in attn.structure_target:
+                value = feature_injection(value, batch_order=attn.batch_order)
+
+        # Appearance transfer (value)
+        if do_appearance_control and "value" in attn.appearance_target:
+            value = appearance_transfer(value, query_normed, key_normed, batch_order=attn.batch_order)
+
+        # The output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        # Appearance transfer (after)
+        if do_appearance_control and "after" in attn.appearance_target:
+            hidden_states = appearance_transfer(hidden_states, query_normed, key_normed, batch_order=attn.batch_order)
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # Linear projection
+        hidden_states = attn.to_out[0](hidden_states, *args)
+        # 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 register_control(
+    model,
+    timesteps,
+    control_schedule,  # structure_conv, structure_attn, appearance_attn
+    control_target = [["output_tensor"], ["query", "key"], ["before"]],
+):
+    # Assume timesteps in reverse order (T -> 0)
+    for block_type in ["encoder", "decoder", "middle"]:
+        blocks = {
+            "encoder": model.unet.down_blocks,
+            "decoder": model.unet.up_blocks,
+            "middle": [model.unet.mid_block],
+        }[block_type]
+
+        control_schedule_block = control_schedule[block_type]
+        if block_type == "middle":
+            control_schedule_block = [control_schedule_block]
+
+        for layer in range(len(control_schedule_block)):
+            # Convolution
+            num_blocks = len(blocks[layer].resnets) if hasattr(blocks[layer], "resnets") else 0
+            for block in range(num_blocks):
+                convolution = blocks[layer].resnets[block]
+                convolution.structure_target = control_target[0]
+                convolution.structure_schedule = get_schedule(
+                    timesteps, get_elem(control_schedule_block[layer][0], block)
+                )
+                convolution.forward = MethodType(convolution_forward, convolution)
+
+            # Self-attention
+            num_blocks = len(blocks[layer].attentions) if hasattr(blocks[layer], "attentions") else 0
+            for block in range(num_blocks):
+                for transformer_block in blocks[layer].attentions[block].transformer_blocks:
+                    attention = transformer_block.attn1
+                    attention.structure_target = control_target[1]
+                    attention.structure_schedule = get_schedule(
+                        timesteps, get_elem(control_schedule_block[layer][1], block)
+                    )
+                    attention.appearance_target = control_target[2]
+                    attention.appearance_schedule = get_schedule(
+                        timesteps, get_elem(control_schedule_block[layer][2], block)
+                    )
+                    attention.processor = AttnProcessor2_0()
+
+
+def register_attr(model, t, do_control, batch_order):
+    for layer_type in ["encoder", "decoder", "middle"]:
+        blocks = {"encoder": model.unet.down_blocks, "decoder": model.unet.up_blocks,
+                  "middle": [model.unet.mid_block]}[layer_type]
+        for layer in blocks:
+            # Convolution
+            for module in layer.resnets:
+                module.t = t
+                module.do_control = do_control
+                module.batch_order = batch_order
+            # Self-attention
+            if hasattr(layer, "attentions"):
+                for block in layer.attentions:
+                    for module in block.transformer_blocks:
+                        module.attn1.t = t
+                        module.attn1.do_control = do_control
+                        module.attn1.batch_order = batch_order

modules/ctrlx/utils.py

@@ -0,0 +1,100 @@
+import random
+from os import environ
+import numpy as np
+import torch
+
+
+JPEG_QUALITY = 100
+
+
+def seed_everything(seed):
+    random.seed(seed)
+    environ["PYTHONHASHSEED"] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def exists(x):
+    return x is not None
+
+
+def get(x, default):
+    if exists(x):
+        return x
+    return default
+
+
+def get_self_recurrence_schedule(schedule, num_inference_steps):
+    self_recurrence_schedule = [0] * num_inference_steps
+    for schedule_current in reversed(schedule):
+        if schedule_current is None or len(schedule_current) == 0:
+            continue
+        [start, end, repeat] = schedule_current
+        start_i = round(num_inference_steps * start)
+        end_i = round(num_inference_steps * end)
+        for i in range(start_i, end_i):
+            self_recurrence_schedule[i] = repeat
+    return self_recurrence_schedule
+
+
+def batch_dict_to_tensor(batch_dict, batch_order):
+    batch_tensor = []
+    for batch_type in batch_order:
+        batch_tensor.append(batch_dict[batch_type])
+    batch_tensor = torch.cat(batch_tensor, dim=0)
+    return batch_tensor
+
+
+def batch_tensor_to_dict(batch_tensor, batch_order):
+    batch_tensor_chunk = batch_tensor.chunk(len(batch_order))
+    batch_dict = {}
+    for i, batch_type in enumerate(batch_order):
+        batch_dict[batch_type] = batch_tensor_chunk[i]
+    return batch_dict
+
+
+def noise_prev(scheduler, timestep, x_0, noise=None):
+    if scheduler.num_inference_steps is None:
+        raise ValueError(
+            "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+        )
+
+    if noise is None:
+        noise = torch.randn_like(x_0).to(x_0)
+
+    # From DDIMScheduler step function (hopefully this works)
+    timestep_i = (scheduler.timesteps == timestep).nonzero(as_tuple=True)[0][0].item()
+    if timestep_i + 1 >= scheduler.timesteps.shape[0]:  # We are at t = 0 (ish)
+        return x_0
+    prev_timestep = scheduler.timesteps[timestep_i + 1:timestep_i + 2]  # Make sure t is not 0-dim
+
+    x_t_prev = scheduler.add_noise(x_0, noise, prev_timestep)
+    return x_t_prev
+
+
+def noise_t2t(scheduler, timestep, timestep_target, x_t, noise=None):
+    assert timestep_target >= timestep
+    if noise is None:
+        noise = torch.randn_like(x_t).to(x_t)
+
+    alphas_cumprod = scheduler.alphas_cumprod.to(device=x_t.device, dtype=x_t.dtype)
+
+    timestep = timestep.to(torch.long)
+    timestep_target = timestep_target.to(torch.long)
+
+    alpha_prod_t = alphas_cumprod[timestep]
+    alpha_prod_tt = alphas_cumprod[timestep_target]
+    alpha_prod = alpha_prod_tt / alpha_prod_t
+
+    sqrt_alpha_prod = (alpha_prod ** 0.5).flatten()
+    while len(sqrt_alpha_prod.shape) < len(x_t.shape):
+        sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+    sqrt_one_minus_alpha_prod = ((1 - alpha_prod) ** 0.5).flatten()
+    while len(sqrt_one_minus_alpha_prod.shape) < len(x_t.shape):
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+    x_tt = sqrt_alpha_prod * x_t + sqrt_one_minus_alpha_prod * noise
+    return x_tt

modules/devices.py

@@ -228,14 +228,7 @@ def test_bf16():
         return False
 
 
-def set_cuda_params():
-    if debug:
-        log.debug(f'Verifying Torch settings: cuda={cuda_ok}')
-    if backend == "ipex":
-        try:
-            torch.xpu.set_fp32_math_mode(mode=torch.xpu.FP32MathMode.TF32)
-        except Exception:
-            pass
+def set_cudnn_params():
     if cuda_ok:
         try:
             torch.backends.cuda.matmul.allow_tf32 = True
@@ -257,6 +250,17 @@ def set_cuda_params():
                 torch.backends.cudnn.allow_tf32 = True
             except Exception:
                 pass
+
+
+def override_ipex_math():
+    if backend == "ipex":
+        try:
+            torch.xpu.set_fp32_math_mode(mode=torch.xpu.FP32MathMode.TF32)
+        except Exception:
+            pass
+
+
+def set_sdpa_params():
     try:
         if shared.opts.cross_attention_optimization == "Scaled-Dot-Product":
             torch.backends.cuda.enable_flash_sdp('Flash attention' in shared.opts.sdp_options)
@@ -283,8 +287,9 @@ def set_cuda_params():
                 torch.nn.functional.scaled_dot_product_attention = sliced_scaled_dot_product_attention
     except Exception:
         pass
-    if shared.cmd_opts.profile:
-        shared.log.debug(f'Torch info: {torch.__config__.show()}')
+
+
+def set_dtype():
     global dtype, dtype_vae, dtype_unet, unet_needs_upcast, inference_context, fp16_ok, bf16_ok # pylint: disable=global-statement
     if shared.opts.cuda_dtype == 'FP32':
         dtype = torch.float32
@@ -319,9 +324,20 @@ def set_cuda_params():
         inference_context = contextlib.nullcontext
     else:
         inference_context = torch.no_grad
-    log_device_name = get_raw_openvino_device() if shared.cmd_opts.use_openvino else torch.device(get_optimal_device_name()) # pylint: disable=used-before-assignment
+
+
+def set_cuda_params():
+    if debug:
+        log.debug(f'Verifying Torch settings: cuda={cuda_ok}')
+    override_ipex_math()
+    set_cudnn_params()
+    set_sdpa_params()
+    set_dtype()
+    if shared.cmd_opts.profile:
+        shared.log.debug(f'Torch info: {torch.__config__.show()}')
+    device_name = get_raw_openvino_device() if shared.cmd_opts.use_openvino else torch.device(get_optimal_device_name()) # pylint: disable=used-before-assignment
     log.debug(f'Desired Torch parameters: dtype={shared.opts.cuda_dtype} no-half={shared.opts.no_half} no-half-vae={shared.opts.no_half_vae} upscast={shared.opts.upcast_sampling}')
-    log.info(f'Setting Torch parameters: device={log_device_name} dtype={dtype} vae={dtype_vae} unet={dtype_unet} context={inference_context.__name__} fp16={fp16_ok} bf16={bf16_ok} optimization={shared.opts.cross_attention_optimization}')
+    log.info(f'Setting Torch parameters: device={device_name} dtype={dtype} vae={dtype_vae} unet={dtype_unet} context={inference_context.__name__} fp16={fp16_ok} bf16={bf16_ok} optimization={shared.opts.cross_attention_optimization}')
 
 
 args = cmd_args.parser.parse_args()

modules/errors.py

@@ -17,7 +17,7 @@ console = Console(log_time=True, tab_size=4, log_time_format='%H:%M:%S-%f', soft
 }))
 
 pretty_install(console=console)
-traceback_install(console=console, extra_lines=1, width=console.width, word_wrap=False, indent_guides=False)
+traceback_install(console=console, extra_lines=1, width=console.width, word_wrap=False, indent_guides=False, max_frames=16)
 already_displayed = {}
 
 
@@ -38,7 +38,7 @@ def print_error_explanation(message):
 
 def display(e: Exception, task, suppress=[]):
     log.error(f"{task or 'error'}: {type(e).__name__}")
-    console.print_exception(show_locals=False, max_frames=10, extra_lines=1, suppress=suppress, theme="ansi_dark", word_wrap=False, width=console.width)
+    console.print_exception(show_locals=False, max_frames=16, extra_lines=1, suppress=suppress, theme="ansi_dark", word_wrap=False, width=console.width)
 
 
 def display_once(e: Exception, task):
@@ -56,7 +56,7 @@ def run(code, task):
 
 
 def exception(suppress=[]):
-    console.print_exception(show_locals=False, max_frames=10, extra_lines=2, suppress=suppress, theme="ansi_dark", word_wrap=False, width=min([console.width, 200]))
+    console.print_exception(show_locals=False, max_frames=16, extra_lines=2, suppress=suppress, theme="ansi_dark", word_wrap=False, width=min([console.width, 200]))
 
 
 def profile(profiler, msg: str):

modules/images_namegen.py

@@ -177,6 +177,7 @@ class FilenameGenerator:
         return fn
 
     def sequence(self, fn, dirname, basename):
+        x = fn
         if shared.opts.save_images_add_number or '[seq]' in fn:
             if '[seq]' not in fn:
                 fn = os.path.join(os.path.dirname(fn), f"[seq]-{os.path.basename(fn)}")

modules/modelloader.py

@@ -36,7 +36,7 @@ def hf_login(token=None):
         stdout = io.StringIO()
         with contextlib.redirect_stdout(stdout):
             hf.logout()
-            hf.login(token=shared.opts.huggingface_token, add_to_git_credential=False, write_permission=False)
+            hf.login(token=token, add_to_git_credential=False, write_permission=False)
         text = stdout.getvalue() or ''
         line = [l for l in text.split('\n') if 'Token' in l]
         shared.log.info(f'HF login: token="{hf.constants.HF_TOKEN_PATH}" {line[0] if len(line) > 0 else text}')

modules/processing_vae.py

@@ -56,7 +56,9 @@ def full_vae_decode(latents, model):
     if upcast:
         if hasattr(model, 'upcast_vae'): # this is done by diffusers automatically if output_type != 'latent'
             model.upcast_vae()
-        model.vae = model.vae.to(dtype=torch.float32)
+        else: # manual upcast and we restore it later
+            model.vae.orig_dtype = model.vae.dtype
+            model.vae = model.vae.to(dtype=torch.float32)
         latents = latents.to(torch.float32)
     else:
         latents = latents.to(devices.device)
@@ -89,6 +91,10 @@ def full_vae_decode(latents, model):
         errors.display(e, 'VAE decode')
         decoded = []
 
+    if hasattr(model.vae, "orig_dtype"):
+        model.vae = model.vae.to(dtype=model.vae.orig_dtype)
+        del model.vae.orig_dtype
+
     # delete vae after OpenVINO compile
     if 'VAE' in shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx" and shared.compiled_model_state.first_pass_vae:
         shared.compiled_model_state.first_pass_vae = False

modules/rocm.py

@@ -24,8 +24,8 @@ def dirname(path_: str, r: int = 1) -> str:
     return path_
 
 
-def spawn(command: str) -> str:
-    process = subprocess.run(command, shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+def spawn(command: str, cwd: os.PathLike = '.') -> str:
+    process = subprocess.run(command, cwd=cwd, shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
     return process.stdout.decode(encoding="utf8", errors="ignore")
 
 
@@ -66,12 +66,14 @@ class Agent:
             self.arch = MicroArchitecture.CDNA
         else:
             self.arch = MicroArchitecture.GCN
-        self.is_apu = gfx in ("801", "902", "90c", "1013", "1033", "1035", "1036", "1103",)
+        self.is_apu = gfx.startswith("115") or gfx in ("801", "902", "90c", "1013", "1033", "1035", "1036", "1103",)
         if sys.platform != "win32":
             self.blaslt_supported = os.path.exists(os.path.join(HIPBLASLT_TENSILE_LIBPATH, f"extop_{name}.co"))
 
     def get_gfx_version(self) -> Union[str, None]:
-        if self.name.startswith("gfx11"):
+        if self.name.startswith("gfx12"):
+            return "12.0.0"
+        elif self.name.startswith("gfx11"):
             return "11.0.0"
         elif self.name.startswith("gfx103"): # or self.name.startswith("gfx101")
             return "10.3.0"
@@ -140,7 +142,7 @@ if sys.platform == "win32":
         return os.path.basename(path) or os.path.basename(os.path.dirname(path))
 
     def get_agents() -> List[Agent]:
-        return [Agent(x.split(' ')[-1].strip()) for x in spawn("hipinfo").split("\n") if x.startswith('gcnArchName:')]
+        return [Agent(x.split(' ')[-1].strip()) for x in spawn("hipinfo", cwd=os.path.join(path, 'bin')).split("\n") if x.startswith('gcnArchName:')]
 
     is_wsl: bool = False
 else:
@@ -153,7 +155,7 @@ else:
         return resolve_link("/opt/rocm")
 
     def get_version() -> str:
-        arr = spawn(f"{os.path.join(path, 'bin', 'hipconfig')} --version").split(".")
+        arr = spawn("hipconfig --version", cwd=os.path.join(path, 'bin')).split(".")
         return f'{arr[0]}.{arr[1]}' if len(arr) >= 2 else None
 
     def get_agents() -> List[Agent]:

modules/sd_models.py

@@ -1382,7 +1382,7 @@ def get_signature(cls):
     return signature.parameters
 
 
-def switch_pipe(cls: diffusers.DiffusionPipeline, pipeline: diffusers.DiffusionPipeline = None, args = {}):
+def switch_pipe(cls: diffusers.DiffusionPipeline, pipeline: diffusers.DiffusionPipeline = None, force = False, args = {}):
     """
     args:
     - cls: can be pipeline class or a string from custom pipelines
@@ -1400,7 +1400,7 @@ def switch_pipe(cls: diffusers.DiffusionPipeline, pipeline: diffusers.DiffusionP
         new_pipe = None
         signature = get_signature(cls)
         possible = signature.keys()
-        if isinstance(pipeline, cls) and args == {}:
+        if not force and isinstance(pipeline, cls) and args == {}:
             return pipeline
         pipe_dict = {}
         components_used = []

modules/sd_samplers.py

@@ -78,7 +78,7 @@ def create_sampler(name, model):
             shared.log.warning(f'AlphaVLLM-Lumina: sampler="{name}" unsupported')
             return None
         if not hasattr(model, 'scheduler_config'):
-            model.scheduler_config = sampler.sampler.config.copy()
+            model.scheduler_config = sampler.sampler.config.copy() if hasattr(sampler.sampler, 'config') else {}
         model.scheduler = sampler.sampler
         if hasattr(model, "prior_pipe") and hasattr(model.prior_pipe, "scheduler"):
             model.prior_pipe.scheduler = sampler.sampler

modules/sd_samplers_diffusers.py

@@ -54,8 +54,9 @@ config = {
     'DPM++ 2M': { 'thresholding': False, 'sample_max_value': 1.0, 'algorithm_type': "dpmsolver++", 'solver_type': "midpoint", 'lower_order_final': True, 'use_karras_sigmas': False, 'final_sigmas_type': 'zero', 'timestep_spacing': 'linspace', 'solver_order': 2 },
     'DPM++ 3M': { 'thresholding': False, 'sample_max_value': 1.0, 'algorithm_type': "dpmsolver++", 'solver_type': "midpoint", 'lower_order_final': True, 'use_karras_sigmas': False, 'final_sigmas_type': 'zero', 'timestep_spacing': 'linspace', 'solver_order': 3 },
     'DPM SDE': { 'use_karras_sigmas': False, 'noise_sampler_seed': None, 'timestep_spacing': 'linspace', 'steps_offset': 0 },
-    'Euler a': { 'steps_offset': 0, 'rescale_betas_zero_snr': False, 'timestep_spacing': 'linspace' },
     'Euler': { 'steps_offset': 0, 'interpolation_type': "linear", 'use_karras_sigmas': False, 'rescale_betas_zero_snr': False, 'final_sigmas_type': 'zero', 'timestep_spacing': 'linspace' },
+    'Euler a': { 'steps_offset': 0, 'rescale_betas_zero_snr': False, 'timestep_spacing': 'linspace' },
+    'Euler SGM': { 'timestep_spacing': "trailing", 'prediction_type': "sample" },
     'Heun': { 'use_karras_sigmas': False, 'timestep_spacing': 'linspace' },
     'DDPM': { 'variance_type': "fixed_small", 'clip_sample': False, 'thresholding': False, 'clip_sample_range': 1.0, 'sample_max_value': 1.0, 'timestep_spacing': 'linspace', 'rescale_betas_zero_snr': False },
     'KDPM2': { 'steps_offset': 0, 'timestep_spacing': 'linspace' },
@@ -66,7 +67,6 @@ config = {
     'DC Solver': { 'beta_start': 0.0001, 'beta_end': 0.02, 'solver_order': 2, 'prediction_type': "epsilon", 'thresholding': False, 'solver_type': 'bh2', 'lower_order_final': True, 'dc_order': 2, 'disable_corrector': [0] },
     'LCM': { 'beta_start': 0.00085, 'beta_end': 0.012, 'beta_schedule': "scaled_linear", 'set_alpha_to_one': True, 'rescale_betas_zero_snr': False, 'thresholding': False, 'timestep_spacing': 'linspace' },
     'TCD': { 'set_alpha_to_one': True, 'rescale_betas_zero_snr': False, 'beta_schedule': 'scaled_linear' },
-    'Euler SGM': { 'timestep_spacing': "trailing", 'prediction_type': "sample" },
     'Euler EDM': { },
     'Variational VDM': { 'clip_sample_range': 2.0, },
     'DPM++ 2M EDM': { 'solver_order': 2, 'solver_type': 'midpoint', 'final_sigmas_type': 'zero', 'algorithm_type': 'dpmsolver++' },
@@ -87,6 +87,7 @@ samplers_data_diffusers = [
     sd_samplers_common.SamplerData('Heun', lambda model: DiffusionSampler('Heun', HeunDiscreteScheduler, model), [], {}),
     sd_samplers_common.SamplerData('Euler', lambda model: DiffusionSampler('Euler', EulerDiscreteScheduler, model), [], {}),
     sd_samplers_common.SamplerData('Euler a', lambda model: DiffusionSampler('Euler a', EulerAncestralDiscreteScheduler, model), [], {}),
+    sd_samplers_common.SamplerData('Euler SGM', lambda model: DiffusionSampler('Euler SGM', EulerDiscreteScheduler, model), [], {}),
     sd_samplers_common.SamplerData('DPM++', lambda model: DiffusionSampler('DPM++', DPMSolverSinglestepScheduler, model), [], {}),
     sd_samplers_common.SamplerData('DPM++ 1S', lambda model: DiffusionSampler('DPM++ 1S', DPMSolverMultistepScheduler, model), [], {}),
     sd_samplers_common.SamplerData('DPM++ 2M', lambda model: DiffusionSampler('DPM++ 2M', DPMSolverMultistepScheduler, model), [], {}),
@@ -101,7 +102,6 @@ samplers_data_diffusers = [
     sd_samplers_common.SamplerData('KDPM2 a', lambda model: DiffusionSampler('KDPM2 a', KDPM2AncestralDiscreteScheduler, model), [], {}),
     sd_samplers_common.SamplerData('DPM++ 2M EDM', lambda model: DiffusionSampler('DPM++ 2M EDM', EDMDPMSolverMultistepScheduler, model), [], {}),
     sd_samplers_common.SamplerData('Euler EDM', lambda model: DiffusionSampler('Euler EDM', EDMEulerScheduler, model), [], {}),
-    sd_samplers_common.SamplerData('Euler SGM', lambda model: DiffusionSampler('Euler SGM', EulerDiscreteScheduler, model), [], {}),
     sd_samplers_common.SamplerData('LCM', lambda model: DiffusionSampler('LCM', LCMScheduler, model), [], {}),
     sd_samplers_common.SamplerData('TCD', lambda model: DiffusionSampler('TCD', TCDScheduler, model), [], {}),
     sd_samplers_common.SamplerData('CMSI', lambda model: DiffusionSampler('CMSI', CMStochasticIterativeScheduler, model), [], {}),

modules/ui_models.py

@@ -391,7 +391,7 @@ def create_ui():
                                 hf_selected = gr.Textbox('', label='Select model', placeholder='select model from search results or enter model name manually')
                         with gr.Column(scale=1):
                             with gr.Row():
-                                hf_variant = gr.Textbox(opts.cuda_dtype.lower(), label='Specify model variant', placeholder='')
+                                hf_variant = gr.Textbox('', label='Specify model variant', placeholder='')
                                 hf_revision = gr.Textbox('', label='Specify model revision', placeholder='')
                     with gr.Row():
                         hf_token = gr.Textbox('', label='Huggingface token', placeholder='optional access token for private or gated models')

modules/vdm/__init__.py

@@ -406,8 +406,8 @@ class VDMScheduler(SchedulerMixin, ConfigMixin):
         sqrt_alpha_prod = torch.sqrt(torch.sigmoid(gamma))
         sqrt_one_minus_alpha_prod = torch.sqrt(torch.sigmoid(-gamma))  # sqrt(sigma)
 
-        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
-        return noisy_samples
+        noisy_samples = original_samples * sqrt_alpha_prod + noise * sqrt_one_minus_alpha_prod
+        return noisy_samples.to(original_samples.dtype)
 
     def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor:
         gamma = self.log_snr(timesteps).to(sample.device)

scripts/ctrlx.py

@@ -0,0 +1,92 @@
+# https://github.com/genforce/ctrl-x
+
+import gradio as gr
+from modules import shared, scripts, processing, processing_helpers, sd_models, devices
+
+
+class Script(scripts.Script):
+    def title(self):
+        return 'Ctrl-X'
+
+    def show(self, is_img2img):
+        return shared.native
+
+    def ui(self, _is_img2img):
+        with gr.Row():
+            gr.HTML('<a href="https://github.com/genforce/ctrl-x">&nbsp Ctrl-X</a><br>')
+        with gr.Accordion(label='Structure', open=True):
+            with gr.Row():
+                struct_prompt = gr.Textbox(label='Prompt', value='', rows=1)
+            with gr.Row():
+                struct_strength = gr.Slider(label='Strength', value=0.5, minimum=0.0, maximum=1.0, step=0.05)
+                struct_guidance = gr.Slider(label='Guidance', value=5.0, minimum=0.0, maximum=14.0, step=0.05)
+            with gr.Row():
+                struct_image = gr.Image(label='Image', source='upload', type='pil')
+        with gr.Accordion(label='Appearance', open=True):
+            with gr.Row():
+                appear_prompt = gr.Textbox(label='Prompt', value='', rows=1)
+            with gr.Row():
+                appear_strength = gr.Slider(label='Strength', value=0.5, minimum=0.0, maximum=1.0, step=0.05)
+                appear_guidance = gr.Slider(label='Guidance', value=5.0, minimum=0.0, maximum=14.0, step=0.05)
+            with gr.Row():
+                appear_image = gr.Image(label='Image', source='upload', type='pil')
+        return struct_prompt, struct_strength, struct_guidance, struct_image, appear_prompt, appear_strength, appear_guidance, appear_image
+
+    def run(self, p: processing.StableDiffusionProcessing, struct_prompt, struct_strength, struct_guidance, struct_image, appear_prompt, appear_strength, appear_guidance, appear_image): # pylint: disable=arguments-differ
+        c = shared.sd_model.__class__.__name__ if shared.sd_loaded else ''
+        if shared.sd_model_type != 'sdxl':
+            shared.log.warning(f'Ctrl-X: pipeline={c} required=StableDiffusionXLPipeline')
+            return None
+
+        import yaml
+        from diffusers import StableDiffusionXLPipeline
+        from modules.ctrlx import CtrlXStableDiffusionXLPipeline
+        from modules.ctrlx.sdxl import get_control_config, register_control
+        from modules.ctrlx.utils import get_self_recurrence_schedule
+
+        orig_prompt_attention = shared.opts.prompt_attention
+        shared.opts.data['prompt_attention'] = 'Fixed attention'
+        shared.sd_model = sd_models.switch_pipe(CtrlXStableDiffusionXLPipeline, shared.sd_model)
+
+        # calculate ctrx+x schedule
+        if p.sampler_name not in ['DDIM', 'Euler', 'Euler a', 'DPM++ 1S', 'DDPM', 'Euler SGM', 'LCM', 'TCD']:
+            shared.log.warning(f'Ctrl-X: sampler={p.sampler_name} override="Euler a" supported=[Euler, Euler a, Euler SGM, DDIM, DDPM, , LCM, TCD]')
+            p.sampler_name = 'Euler a'
+        processing_helpers.update_sampler(p, shared.sd_model)
+        shared.sd_model.scheduler.set_timesteps(p.steps, device=devices.device)
+        timesteps = shared.sd_model.scheduler.timesteps
+        control_config = get_control_config(structure_schedule=struct_strength, appearance_schedule=appear_strength)
+        config = yaml.safe_load(control_config)
+        register_control(
+            model=shared.sd_model,
+            timesteps=timesteps,
+            control_schedule=config['control_schedule'],
+            control_target=config['control_target'],
+        )
+
+        # set args
+        if struct_image is not None:
+            p.task_args['structure_prompt'] = struct_prompt
+            p.task_args['structure_image'] = struct_image
+            p.task_args['structure_guidance_scale'] = struct_guidance
+        if appear_image is not None:
+            p.task_args['appearance_prompt'] = appear_prompt
+            p.task_args['appearance_image'] = appear_image
+            p.task_args['appearance_guidance_scale'] = appear_guidance
+        elif hasattr(p, 'init_images') and p.init_images is not None and len(p.init_images) > 0:
+            p.task_args['appearance_image'] = p.init_images[0]
+            p.init_images = None
+        p.task_args['control_schedule'] = config['control_schedule']
+        p.task_args['self_recurrence_schedule'] = get_self_recurrence_schedule(config['self_recurrence_schedule'], p.steps)
+        is_struct = p.task_args.get('structure_image') is not None
+        is_appear = p.task_args.get('appearance_image') is not None
+        shared.log.info(f'Ctrl-X: structure={struct_strength if is_struct else None} appearance={appear_strength if is_appear else None}')
+        shared.log.debug(f'Ctrl-X: config={control_config} args={p.task_args}')
+
+        # process
+        processed: processing.Processed = processing.process_images(p)
+
+        # restore and return
+        shared.opts.data['prompt_attention'] = orig_prompt_attention
+        shared.sd_model = sd_models.switch_pipe(StableDiffusionXLPipeline, shared.sd_model, force=True)
+        return processed