uploading remaining files

model_index.json

@@ -0,0 +1,10 @@
+{
+  "_class_name": "SuperDiffPipeline",
+  "_diffusers_version": "0.31.0",
+  "batch_size": null,
+  "device": "cuda",
+  "guidance_scale": null,
+  "lift": null,
+  "num_inference_steps": null,
+  "seed": null
+}

pipeline.py

@@ -0,0 +1,375 @@
+import random
+from typing import Callable, Dict, List, Optional
+
+import torch
+from diffusers import DiffusionPipeline
+from diffusers.configuration_utils import ConfigMixin
+
+
+class SuperDiffPipeline(DiffusionPipeline, ConfigMixin):
+    """SuperDiffPipeline."""
+
+    def __init__(self, model: Callable, vae: Callable, text_encoder: Callable, scheduler: Callable, tokenizer: Callable, **kwargs) -> 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__()
+        self.model = model
+        self.vae = vae
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.scheduler = scheduler
+
+        device = "cuda" if torch.cuda.is_available() else "cpu" 
+
+        self.vae.to(device)
+        self.model.to(device)
+        self.text_encoder.to(device)
+
+        self.register_to_config(
+            #model=model,
+            #vae=vae,
+            #tokenizer=tokenizer,
+            #text_encoder=text_encoder,
+            #scheduler=scheduler,
+            device=device,
+            batch_size=None,
+            num_inference_steps=None,
+            guidance_scale=None,
+            lift=None,
+            seed=None,
+        )
+
+    @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(
+            _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 = 1000,
+        batch_size: int = 1,
+        lift: int = 0.0,
+        height: int = 512,
+        width: int = 512,
+        guidance_scale: int = 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
+        lift : int
+            lift
+        height : int
+            height
+        width : int
+            width
+        guidance_scale : int
+            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.lift = lift
+        self.seed = seed
+        if self.seed is None:
+            self.seed = random.randint(0, 2**32 - 1)
+        obj_prompt = [prompt_1]
+        bg_prompt = [prompt_2]
+        obj_embeddings = self.get_text_embedding(obj_prompt * batch_size)
+        bg_embeddings = self.get_text_embedding(bg_prompt * batch_size)
+
+        uncond_embeddings = self.get_text_embedding([""] * batch_size)
+
+        generator = torch.cuda.manual_seed(
+            self.seed
+        )  # Seed generator to create the initial latent noise
+        latents = torch.randn(
+            (batch_size, self.model.config.in_channels, height // 8, width // 8),
+            generator=generator,
+            device=self.device,
+        )
+
+        latents_og = latents.clone().detach()
+        latents_uncond_og = latents.clone().detach()
+
+        self.scheduler.set_timesteps(num_inference_steps)
+        latents = latents * self.scheduler.init_noise_sigma
+
+        latents_uncond = latents.clone().detach()
+        return {
+            "latents": latents,
+            "obj_embeddings": obj_embeddings,
+            "uncond_embeddings": uncond_embeddings,
+            "bg_embeddings": bg_embeddings,
+        }
+
+    def _forward(self, model_inputs: Dict) -> Callable:
+        """_forward.
+
+        Parameters
+        ----------
+        model_inputs : Dict
+            model_inputs
+
+        Returns
+        -------
+        Callable
+
+        """
+        latents = model_inputs["latents"]
+        obj_embeddings = model_inputs["obj_embeddings"]
+        uncond_embeddings = model_inputs["uncond_embeddings"]
+        bg_embeddings = model_inputs["bg_embeddings"]
+
+        kappa = 0.5 * torch.ones(
+            (self.num_inference_steps + 1, self.batch_size), device=self.device
+        )
+        ll_obj = torch.ones(
+            (self.num_inference_steps + 1, self.batch_size), device=self.device
+        )
+        ll_bg = torch.ones(
+            (self.num_inference_steps + 1, self.batch_size), device=self.device
+        )
+        ll_uncond = torch.ones(
+            (self.num_inference_steps + 1, self.batch_size), device=self.device
+        )
+        with torch.no_grad():
+            for i, t in enumerate(self.scheduler.timesteps):
+                dsigma = self.scheduler.sigmas[i +
+                                               1] - self.scheduler.sigmas[i]
+                sigma = self.scheduler.sigmas[i]
+                vel_obj = self.get_vel(t, sigma, latents, [obj_embeddings])
+                vel_uncond = self.get_vel(
+                    t, sigma, latents, [uncond_embeddings])
+
+                vel_bg = self.get_vel(t, sigma, latents, [bg_embeddings])
+                noise = torch.sqrt(2 * torch.abs(dsigma) * sigma) * torch.randn_like(
+                    latents
+                )
+
+                dx_ind = (
+                    2
+                    * dsigma
+                    * (vel_uncond + self.guidance_scale * (vel_bg - vel_uncond))
+                    + noise
+                )
+                kappa[i + 1] = (
+                    (torch.abs(dsigma) * (vel_bg - vel_obj) * (vel_bg + vel_obj)).sum(
+                        (1, 2, 3)
+                    )
+                    - (dx_ind * ((vel_obj - vel_bg))).sum((1, 2, 3))
+                    + sigma * self.lift / self.num_inference_steps
+                )
+                kappa[i + 1] /= (
+                    2
+                    * dsigma
+                    * self.guidance_scale
+                    * ((vel_obj - vel_bg) ** 2).sum((1, 2, 3))
+                )
+
+                vf = vel_uncond + self.guidance_scale * (
+                    (vel_bg - vel_uncond)
+                    + kappa[i + 1][:, None, None, None] * (vel_obj - vel_bg)
+                )
+                dx = 2 * dsigma * vf + noise
+                latents += dx
+
+                ll_obj[i + 1] = ll_obj[i] + (
+                    -torch.abs(dsigma) / sigma * (vel_obj) ** 2
+                    - (dx * (vel_obj / sigma))
+                ).sum((1, 2, 3))
+                ll_bg[i + 1] = ll_bg[i] + (
+                    -torch.abs(dsigma) / sigma * (vel_bg) ** 2 -
+                    (dx * (vel_bg / sigma))
+                ).sum((1, 2, 3))
+
+        return latents
+
+    def postprocess(self, latents: Callable) -> Callable:
+        """postprocess.
+
+        Parameters
+        ----------
+        latents : Callable
+            latents
+
+        Returns
+        -------
+        Callable
+
+        """
+        image = self.get_batch(latents, 1, self.batch_size)
+        # Ensure the shape is (height, width, 3)
+        assert image.shape[-1] == 3  # Handle grayscale or invalid shapes
+
+        # Convert to uint8 if not already
+        image = image.to(torch.uint8)  # Ensure it's uint8 for PIL
+
+        return image
+
+    def __call__(
+        self,
+        prompt_1: str,
+        prompt_2: str,
+        seed: int = None,
+        num_inference_steps: int = 1000,
+        batch_size: int = 1,
+        lift: int = 0.0,
+        height: int = 512,
+        width: int = 512,
+        guidance_scale: int = 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
+        lift : int
+            lift
+        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,
+            lift,
+            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

scheduler/scheduler_config.json

@@ -0,0 +1,24 @@
+{
+  "_class_name": "EulerDiscreteScheduler",
+  "_diffusers_version": "0.31.0",
+  "beta_end": 0.012,
+  "beta_schedule": "scaled_linear",
+  "beta_start": 0.00085,
+  "clip_sample": false,
+  "final_sigmas_type": "zero",
+  "interpolation_type": "linear",
+  "num_train_timesteps": 1000,
+  "prediction_type": "epsilon",
+  "rescale_betas_zero_snr": false,
+  "set_alpha_to_one": false,
+  "sigma_max": null,
+  "sigma_min": null,
+  "skip_prk_steps": true,
+  "steps_offset": 1,
+  "timestep_spacing": "linspace",
+  "timestep_type": "discrete",
+  "trained_betas": null,
+  "use_beta_sigmas": false,
+  "use_exponential_sigmas": false,
+  "use_karras_sigmas": false
+}

text_encoder/config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "CompVis/stable-diffusion-v1-4",
+  "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": 512,
+  "torch_dtype": "float32",
+  "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:778d02eb9e707c3fbaae0b67b79ea0d1399b52e624fb634f2f19375ae7c047c3
+size 492265168

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": false,
+  "do_lower_case": true,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "model_max_length": 77,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "CLIPTokenizer",
+  "unk_token": "<|endoftext|>"
+}

vae/config.json

@@ -0,0 +1,38 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.31.0",
+  "_name_or_path": "CompVis/stable-diffusion-v1-4",
+  "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": 512,
+  "scaling_factor": 0.18215,
+  "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:b4d2b5932bb4151e54e694fd31ccf51fca908223c9485bd56cd0e1d83ad94c49
+size 334643268