Upload 18 files

model_index.json

@@ -0,0 +1,28 @@
+{
+  "_class_name": "SuperDiffSDXLPipeline",
+  "_diffusers_version": "0.31.0",
+  "text_encoder": [
+    "transformers",
+    "CLIPTextModel"
+  ],
+  "text_encoder_2": [
+    "transformers",
+    "CLIPTextModelWithProjection"
+  ],
+  "tokenizer": [
+    "transformers",
+    "CLIPTokenizer"
+  ],
+  "tokenizer_2": [
+    "transformers",
+    "CLIPTokenizer"
+  ],
+  "unet": [
+    "diffusers",
+    "UNet2DConditionModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ]
+}

pipeline.py

@@ -0,0 +1,417 @@
+import random
+from typing import Callable, Dict, List, Optional
+
+import torch
+from diffusers import DiffusionPipeline
+from diffusers.configuration_utils import ConfigMixin
+from tqdm import tqdm
+
+# from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
+# from diffusers import AutoencoderKL, UNet2DConditionModel
+
+
+def get_scaled_coeffs():
+    """get_scaled_coeffs.
+    """
+    beta_min = 0.85
+    beta_max = 12.0
+    return beta_min**0.5, beta_max**0.5-beta_min**0.5
+
+
+def beta(t):
+    """beta.
+
+    Parameters
+    ----------
+    t :
+        t
+    """
+    a, b = get_scaled_coeffs()
+    return (a+t*b)**2
+
+
+def int_beta(t):
+    """int_beta.
+
+    Parameters
+    ----------
+    t :
+        t
+    """
+  a, b = get_scaled_coeffs()
+  return ((a+b*t)**3-a**3)/(3*b)
+def sigma(t):
+    """sigma.
+
+    Parameters
+    ----------
+    t :
+        t
+    """
+  return torch.expm1(int_beta(t))**0.5
+def sigma_orig(t):
+    """sigma_orig.
+
+    Parameters
+    ----------
+    t :
+        t
+    """
+  return (-torch.expm1(-int_beta(t)))**0.5
+
+class SuperDiffSDXLPipeline(DiffusionPipeline, ConfigMixin):
+    """SuperDiffSDXLPipeline."""
+
+    def __init__(self, unet: Callable, vae: Callable, text_encoder: Callable, text_encoder_2: Callable, tokenizer: Callable, tokenizer_2: Callable) -> None:
+
+        """__init__.
+
+        Parameters
+        ----------
+        model : Callable
+            model
+        vae : Callable
+            vae
+        text_encoder : Callable
+            text_encoder
+        scheduler : Callable
+            scheduler
+        tokenizer : Callable
+            tokenizer
+        kwargs :
+            kwargs
+
+        Returns
+        -------
+        None
+
+        """
+        super().__init__()
+        device = "cuda" if torch.cuda.is_available() else "cpu" 
+        dtype=torch.float16
+
+        vae.to(device)
+        unet.to(device)
+        text_encoder.to(device)
+        text_encoder_2.to(device)
+
+        self.register_modules(unet=unet, 
+                              vae=vae, 
+                              text_encoder=text_encoder,
+                              text_encoder_2=text_encoder_2,
+                              tokenizer=tokenizer,
+                              tokenizer_2=tokenizer_2,
+                              )
+
+    def prepare_prompt_input(self, prompt_o, prompt_b, batch_size, height, width):
+        """prepare_prompt_input.
+
+        Parameters
+        ----------
+        prompt_o :
+            prompt_o
+        prompt_b :
+            prompt_b
+        batch_size :
+            batch_size
+        height :
+            height
+        width :
+            width
+        """
+        text_input = self.tokenizer(prompt_o* batch_size, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt")
+        text_input_2 = self.tokenizer_2(prompt_o* batch_size, padding="max_length", max_length=self.tokenizer_2.model_max_length, truncation=True, return_tensors="pt")
+        with torch.no_grad():
+            text_embeddings = self.text_encoder(text_input.input_ids.to(self.device), output_hidden_states=True)
+            text_embeddings_2 = self.text_encoder_2(text_input_2.input_ids.to(self.device), output_hidden_states=True)
+        prompt_embeds_o = torch.concat((text_embeddings.hidden_states[-2], text_embeddings_2.hidden_states[-2]), dim=-1)
+        pooled_prompt_embeds_o = text_embeddings_2[0]
+        negative_prompt_embeds = torch.zeros_like(prompt_embeds_o)
+        negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds_o)
+    
+        text_input = self.tokenizer(prompt_b* batch_size, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt")
+        text_input_2 = self.tokenizer_2(prompt_b* batch_size, padding="max_length", max_length=self.tokenizer_2.model_max_length, truncation=True, return_tensors="pt")
+        with torch.no_grad():
+            text_embeddings = self.text_encoder(text_input.input_ids.to(self.device), output_hidden_states=True)
+            text_embeddings_2 = self.text_encoder_2(text_input_2.input_ids.to(self.device), output_hidden_states=True)
+        prompt_embeds_b = torch.concat((text_embeddings.hidden_states[-2], text_embeddings_2.hidden_states[-2]), dim=-1)
+        pooled_prompt_embeds_b = text_embeddings_2[0]
+        add_time_ids_o = torch.tensor([(height,width,0,0,height,width)])
+        add_time_ids_b = torch.tensor([(height,width,0,0,height,width)])
+        negative_add_time_ids = torch.tensor([(height,width,0,0,height,width)])
+        prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds_o, prompt_embeds_b], dim=0)
+        add_text_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds_o, pooled_prompt_embeds_b], dim=0)
+        add_time_ids = torch.cat([negative_add_time_ids, add_time_ids_o, add_time_ids_b], dim=0)
+    
+        prompt_embeds = prompt_embeds.to(self.device)
+        add_text_embeds = add_text_embeds.to(self.device)
+        add_time_ids = add_time_ids.to(self.device).repeat(batch_size, 1)
+        added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
+        return prompt_embeds, added_cond_kwargs
+
+    @torch.no_grad
+    def get_batch(self, latents: Callable, nrow: int, ncol: int) -> Callable:
+        """get_batch.
+
+        Parameters
+        ----------
+        latents : Callable
+            latents
+        nrow : int
+            nrow
+        ncol : int
+            ncol
+
+        Returns
+        -------
+        Callable
+
+        """
+        image = self.vae.decode(
+            latents / self.vae.config.scaling_factor, return_dict=False
+        )[0]
+        image = (image / 2 + 0.5).clamp(0, 1).squeeze()
+        if len(image.shape) < 4:
+            image = image.unsqueeze(0)
+        image = (image.permute(0, 2, 3, 1) * 255).to(torch.uint8)
+        return image
+
+    @torch.no_grad
+    def get_text_embedding(self, prompt: str) -> Callable:
+        """get_text_embedding.
+
+        Parameters
+        ----------
+        prompt : str
+            prompt
+
+        Returns
+        -------
+        Callable
+
+        """
+        text_input = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        return self.text_encoder(text_input.input_ids.to(self.device))[0]
+
+    @torch.no_grad
+    def get_vel(self, t: float, sigma: float, latents: Callable, embeddings: Callable):
+        """get_vel.
+
+        Parameters
+        ----------
+        t : float
+            t
+        sigma : float
+            sigma
+        latents : Callable
+            latents
+        embeddings : Callable
+            embeddings
+        """
+        def v(_x, _e): return self.model(
+            """v.
+
+            Parameters
+            ----------
+            _x :
+                _x
+            _e :
+                _e
+            """
+            _x / ((sigma**2 + 1) ** 0.5), t, encoder_hidden_states=_e
+        ).sample
+        embeds = torch.cat(embeddings)
+        latent_input = latents
+        vel = v(latent_input, embeds)
+        return vel
+
+    def preprocess(
+        self,
+        prompt_1: str,
+        prompt_2: str,
+        seed: int = None,
+        num_inference_steps: int = 200,
+        batch_size: int = 1,
+        height: int = 1024,
+        width: int = 1024,
+        guidance_scale: float = 7.5,
+    ) -> Callable:
+        """preprocess.
+
+        Parameters
+        ----------
+        prompt_1 : str
+            prompt_1
+        prompt_2 : str
+            prompt_2
+        seed : int
+            seed
+        num_inference_steps : int
+            num_inference_steps
+        batch_size : int
+            batch_size
+        height : int
+            height
+        width : int
+            width
+        guidance_scale : float
+            guidance_scale
+
+        Returns
+        -------
+        Callable
+
+        """
+        # Tokenize the input
+        self.batch_size = batch_size
+        self.num_inference_steps = num_inference_steps
+        self.guidance_scale = guidance_scale
+        self.seed = seed
+        if self.seed is None:
+            self.seed = random.randint(0, 2**32 - 1)
+
+        self.generator = torch.cuda.manual_seed(
+            self.seed
+        )  # Seed generator to create the initial latent noise
+
+        latents = torch.randn((batch_size, self.unet.in_channels, height // 8, width // 8), generator=self.generator, dtype=self.dtype, device=self.device,)
+        prompt_embeds, added_cond_kwargs = self.prepare_prompt_input(prompt_1, prompt_2, batch_size, height, width)
+
+        return {
+            "latents": latents,
+            "prompt_embeds": prompt_embeds,
+            "added_cond_kwargs": added_cond_kwargs,
+        }
+
+    def _forward(self, model_inputs: Dict) -> Callable:
+        """_forward.
+
+        Parameters
+        ----------
+        model_inputs : Dict
+            model_inputs
+
+        Returns
+        -------
+        Callable
+
+        """
+        latents = model_inputs["latents"]
+        prompt_embeds = model_inputs["prompt_embeds"]
+        added_cond_kwargs = model_inputs["added_cond_kwargs"]
+
+        t = torch.tensor(1.0)
+        dt = 1.0/self.num_inference_steps
+        train_number_steps = 1000
+        latents = latents * (sigma(t)**2+1)**0.5
+        with torch.no_grad():
+            for i in tqdm(range(self.num_inference_steps)):
+                latent_model_input = torch.cat([latents] * 3)
+                sigma_t = sigma(t)
+                dsigma = sigma(t-dt) - sigma_t
+                latent_model_input /= (sigma_t**2+1)**0.5
+                with torch.no_grad():
+                    noise_pred = self.unet(latent_model_input, t*train_number_steps, encoder_hidden_states=prompt_embeds, added_cond_kwargs=added_cond_kwargs, return_dict=False)[0]
+            
+                noise_pred_uncond, noise_pred_text_o, noise_pred_text_b = noise_pred.chunk(3)
+            
+                # noise = torch.sqrt(2*torch.abs(dsigma)*sigma_t)*torch.randn_like(latents)
+                noise = torch.sqrt(2*torch.abs(dsigma)*sigma_t)*torch.empty_like(latents, device=self.device).normal_(generator=self.generator)
+            
+                dx_ind = 2*dsigma*(noise_pred_uncond + self.guidance_scale*(noise_pred_text_b - noise_pred_uncond)) + noise
+                kappa = (torch.abs(dsigma)*(noise_pred_text_b-noise_pred_text_o)*(noise_pred_text_b+noise_pred_text_o)).sum((1,2,3))-(dx_ind*((noise_pred_text_o-noise_pred_text_b))).sum((1,2,3))
+                kappa /= 2*dsigma*self.guidance_scale*((noise_pred_text_o-noise_pred_text_b)**2).sum((1,2,3))
+                noise_pred = noise_pred_uncond + self.guidance_scale*((noise_pred_text_b - noise_pred_uncond) + kappa[:,None,None,None]*(noise_pred_text_o-noise_pred_text_b))
+            
+                if i < self.num_inference_steps - 1:
+                    latents += 2*dsigma * noise_pred + noise
+                else:
+                    latents += dsigma * noise_pred
+
+                t -= dt
+            return latents
+
+    def postprocess(self, latents: Callable) -> Callable:
+        """postprocess.
+
+        Parameters
+        ----------
+        latents : Callable
+            latents
+
+        Returns
+        -------
+        Callable
+
+        """
+        latents = latents/self.vae.config.scaling_factor
+        latents = latents.to(torch.float32)
+        with torch.no_grad():
+            image = self.vae.decode(latents, return_dict=False)[0]
+        
+        image = (image / 2 + 0.5).clamp(0, 1)
+        image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
+        images = (image * 255).round().astype("uint8")
+        return images
+
+    def __call__(
+        self,
+        prompt_1: str,
+        prompt_2: str,
+        seed: int = None,
+        num_inference_steps: int = 200,
+        batch_size: int = 1,
+        height: int = 1024,
+        width: int = 1024,
+        guidance_scale: float = 7.5,
+    ) -> Callable:
+        """__call__.
+
+        Parameters
+        ----------
+        prompt_1 : str
+            prompt_1
+        prompt_2 : str
+            prompt_2
+        seed : int
+            seed
+        num_inference_steps : int
+            num_inference_steps
+        batch_size : int
+            batch_size
+        height : int
+            height
+        width : int
+            width
+        guidance_scale : int
+            guidance_scale
+
+        Returns
+        -------
+        Callable
+
+        """
+        # Preprocess inputs
+        model_inputs = self.preprocess(
+            prompt_1,
+            prompt_2,
+            seed,
+            num_inference_steps,
+            batch_size,
+            height,
+            width,
+            guidance_scale,
+        )
+
+        # Forward pass through the pipeline
+        latents = self._forward(model_inputs)
+
+        # Postprocess to generate the final output
+        images = self.postprocess(latents)
+        return images

text_encoder/config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "stabilityai/stable-diffusion-xl-base-1.0",
+  "architectures": [
+    "CLIPTextModel"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "dropout": 0.0,
+  "eos_token_id": 2,
+  "hidden_act": "quick_gelu",
+  "hidden_size": 768,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 77,
+  "model_type": "clip_text_model",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 1,
+  "projection_dim": 768,
+  "torch_dtype": "float16",
+  "transformers_version": "4.46.2",
+  "vocab_size": 49408
+}

text_encoder/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:660c6f5b1abae9dc498ac2d21e1347d2abdb0cf6c0c0c8576cd796491d9a6cdd
+size 246144152

text_encoder_2/config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "stabilityai/stable-diffusion-xl-base-1.0",
+  "architectures": [
+    "CLIPTextModelWithProjection"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "dropout": 0.0,
+  "eos_token_id": 2,
+  "hidden_act": "gelu",
+  "hidden_size": 1280,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 5120,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 77,
+  "model_type": "clip_text_model",
+  "num_attention_heads": 20,
+  "num_hidden_layers": 32,
+  "pad_token_id": 1,
+  "projection_dim": 1280,
+  "torch_dtype": "float16",
+  "transformers_version": "4.46.2",
+  "vocab_size": 49408
+}

text_encoder_2/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec310df2af79c318e24d20511b601a591ca8cd4f1fce1d8dff822a356bcdb1f4
+size 1389382176

tokenizer/special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<|endoftext|>",
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer/tokenizer_config.json

@@ -0,0 +1,30 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "49406": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49407": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "do_lower_case": true,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "model_max_length": 77,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "CLIPTokenizer",
+  "unk_token": "<|endoftext|>"
+}

tokenizer_2/special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "!",
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_2/tokenizer_config.json

@@ -0,0 +1,38 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "!",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49406": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49407": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "do_lower_case": true,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "model_max_length": 77,
+  "pad_token": "!",
+  "tokenizer_class": "CLIPTokenizer",
+  "unk_token": "<|endoftext|>"
+}

unet/config.json

@@ -0,0 +1,73 @@
+{
+  "_class_name": "UNet2DConditionModel",
+  "_diffusers_version": "0.31.0",
+  "_name_or_path": "stabilityai/stable-diffusion-xl-base-1.0",
+  "act_fn": "silu",
+  "addition_embed_type": "text_time",
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": 256,
+  "attention_head_dim": [
+    5,
+    10,
+    20
+  ],
+  "attention_type": "default",
+  "block_out_channels": [
+    320,
+    640,
+    1280
+  ],
+  "center_input_sample": false,
+  "class_embed_type": null,
+  "class_embeddings_concat": false,
+  "conv_in_kernel": 3,
+  "conv_out_kernel": 3,
+  "cross_attention_dim": 2048,
+  "cross_attention_norm": null,
+  "down_block_types": [
+    "DownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "dropout": 0.0,
+  "dual_cross_attention": false,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_only_cross_attention": null,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DCrossAttn",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "out_channels": 4,
+  "projection_class_embeddings_input_dim": 2816,
+  "resnet_out_scale_factor": 1.0,
+  "resnet_skip_time_act": false,
+  "resnet_time_scale_shift": "default",
+  "reverse_transformer_layers_per_block": null,
+  "sample_size": 128,
+  "time_cond_proj_dim": null,
+  "time_embedding_act_fn": null,
+  "time_embedding_dim": null,
+  "time_embedding_type": "positional",
+  "timestep_post_act": null,
+  "transformer_layers_per_block": [
+    1,
+    2,
+    10
+  ],
+  "up_block_types": [
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D",
+    "UpBlock2D"
+  ],
+  "upcast_attention": null,
+  "use_linear_projection": true
+}

unet/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:83e012a805b84c7ca28e5646747c90a243c65c8ba4f070e2d7ddc9d74661e139
+size 5135149760

vae/config.json

@@ -0,0 +1,38 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.31.0",
+  "_name_or_path": "stabilityai/stable-diffusion-xl-base-1.0",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 4,
+  "latents_mean": null,
+  "latents_std": null,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 1024,
+  "scaling_factor": 0.13025,
+  "shift_factor": null,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true
+}

vae/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1598f3d24932bcfe6634e8b618ea1e30ab1d57f5aad13a6d2de446d2199f2341
+size 334643268