init space

.gitignore

@@ -0,0 +1,5 @@
+__pycache__
+output_demo/
+cache/
+gradio_cached_examples
+secrets.sh

README.md

@@ -1,12 +1,13 @@
 ---
 title: Diffusion Cocktail
-emoji: 📚
+emoji: 🍸
 colorFrom: indigo
 colorTo: indigo
 sdk: gradio
 sdk_version: 4.9.0
 app_file: app.py
 pinned: false
+python: 3.8
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

ditail/__init__.py

@@ -0,0 +1 @@
+from .src.ditail_demo import DitailDemo, seed_everything

ditail/lora/animeoutline.safetensors

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

ditail/lora/film.safetensors

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

ditail/lora/flat.safetensors

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

ditail/lora/impressionism.safetensors

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

ditail/lora/minecraft.safetensors

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

ditail/lora/pop.safetensors

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

ditail/lora/shinkai_makoto.safetensors

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

ditail/lora/snow.safetensors

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

ditail/src/ditail_demo.py

@@ -0,0 +1,233 @@
+import os
+import yaml
+import argparse
+import warnings
+from PIL import Image
+from tqdm import tqdm
+from datetime import datetime
+
+import torch
+import torch.nn as nn
+import torchvision.transforms as T
+from transformers import logging
+from diffusers import DDIMScheduler, StableDiffusionPipeline
+
+from .ditail_utils import *
+
+# suppress warnings
+logging.set_verbosity_error()
+warnings.filterwarnings("ignore", message=".*LoRA backend.*")
+
+class DitailDemo(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        if isinstance(self.args, dict):
+            for k, v in args.items():
+                setattr(self, k, v)
+        else:
+            for k, v in vars(args).items():
+                setattr(self, k, v)
+
+    def load_inv_model(self):
+        self.scheduler = DDIMScheduler.from_pretrained(self.inv_model, subfolder='scheduler')
+        self.scheduler.set_timesteps(self.inv_steps, device=self.device)
+        print(f'[INFO] Loading inversion model: {self.inv_model}')
+        pipe = StableDiffusionPipeline.from_pretrained(
+            self.inv_model, torch_dtype=torch.float16
+        ).to(self.device)
+        pipe.enable_xformers_memory_efficient_attention()
+        self.text_encoder = pipe.text_encoder
+        self.tokenizer = pipe.tokenizer
+        self.unet = pipe.unet
+        self.vae = pipe.vae
+        self.tokenizer_kwargs = dict(
+            truncation=True,
+            return_tensors='pt',
+            padding='max_length',
+            max_length=self.tokenizer.model_max_length
+        )
+    
+    def load_spl_model(self):
+        self.scheduler = DDIMScheduler.from_pretrained(self.spl_model, subfolder='scheduler')
+        self.scheduler.set_timesteps(self.spl_steps, device=self.device)
+        print(f'[INFO] Loading sampling model: {self.spl_model}')
+        if (self.lora != 'none') or (self.inv_model != self.spl_model):
+            pipe = StableDiffusionPipeline.from_pretrained(
+                self.spl_model, torch_dtype=torch.float16
+            ).to(self.device)
+            if self.lora != 'none':
+                pipe.unfuse_lora()
+                pipe.unload_lora_weights()
+                pipe.load_lora_weights(self.lora_dir, weight_name=f'{self.lora}.safetensors')
+                pipe.fuse_lora(lora_scale=self.lora_scale)
+            pipe.enable_xformers_memory_efficient_attention()
+            self.text_encoder = pipe.text_encoder
+            self.tokenizer = pipe.tokenizer
+            self.unet = pipe.unet
+            self.vae = pipe.vae
+            self.tokenizer_kwargs = dict(
+                truncation=True,
+                return_tensors='pt',
+                padding='max_length',
+                max_length=self.tokenizer.model_max_length
+            )
+
+    @torch.no_grad()
+    def encode_image(self, image_pil):
+        # image_pil = T.Resize(512)(img.convert('RGB'))
+        image_pil = T.Resize(512)(image_pil)
+        image = T.ToTensor()(image_pil).unsqueeze(0).to(self.device)
+        with torch.autocast(device_type=self.device, dtype=torch.float32):
+            image = 2 * image - 1
+            posterior = self.vae.encode(image).latent_dist
+            latent = posterior.mean * 0.18215
+        return latent
+
+    @torch.no_grad()
+    def invert_image(self, cond, latent):
+        self.latents = {}
+        timesteps = reversed(self.scheduler.timesteps)
+        with torch.autocast(device_type=self.device, dtype=torch.float32):
+            for i, t in enumerate(tqdm(timesteps)):
+                cond_batch = cond.repeat(latent.shape[0], 1, 1)
+                alpha_prod_t = self.scheduler.alphas_cumprod[t]
+                alpha_prod_t_prev = (
+                    self.scheduler.alphas_cumprod[timesteps[i-1]]
+                    if i > 0 else self.scheduler.final_alpha_cumprod
+                )
+                mu = alpha_prod_t ** 0.5
+                mu_prev = alpha_prod_t_prev ** 0.5
+                sigma = (1 - alpha_prod_t) ** 0.5
+                sigma_prev = (1 - alpha_prod_t_prev) ** 0.5
+                eps = self.unet(latent, t, encoder_hidden_states=cond_batch).sample
+                pred_x0 = (latent - sigma_prev * eps) / mu_prev
+                latent = mu * pred_x0 + sigma * eps
+                self.latents[t.item()] = latent
+        self.noisy_latent = latent
+
+    @torch.no_grad()
+    def extract_latents(self):
+        # get the embeddings for pos & neg prompts
+        # self.pos_prompt = ' ,'.join(LORA_TRIGGER_WORD.get(self.lora, [''])+[self.pos_prompt])
+        # print('the prompt after adding trigger word:', self.pos_prompt)
+        text_pos = self.tokenizer(self.pos_prompt, **self.tokenizer_kwargs)
+        text_neg = self.tokenizer(self.neg_prompt, **self.tokenizer_kwargs)
+        self.emb_pos = self.text_encoder(text_pos.input_ids.to(self.device))[0]
+        self.emb_neg = self.text_encoder(text_neg.input_ids.to(self.device))[0]
+        # apply condition scaling
+        cond = self.alpha * self.emb_pos - self.beta * self.emb_neg
+        # encode source image & apply DDIM inversion
+        self.invert_image(cond, self.encode_image(self.img))
+
+    @torch.no_grad()
+    def latent_to_image(self, latent, save_path=None):
+        with torch.autocast(device_type=self.device, dtype=torch.float32):
+            latent = 1 / 0.18215 * latent
+            image = self.vae.decode(latent).sample[0]
+            image = (image / 2 + 0.5).clamp(0, 1)
+        # T.ToPILImage()(image).save(save_path)
+        return T.ToPILImage()(image)
+
+    def init_injection(self, attn_ratio=0.5, conv_ratio=0.8):
+        attn_thresh = int(attn_ratio * self.spl_steps)
+        conv_thresh = int(conv_ratio * self.spl_steps)
+        self.attn_inj_timesteps = self.scheduler.timesteps[:attn_thresh]
+        self.conv_inj_timesteps = self.scheduler.timesteps[:conv_thresh]
+        register_attn_inj(self, self.attn_inj_timesteps)
+        register_conv_inj(self, self.conv_inj_timesteps)
+
+    @torch.no_grad()
+    def sampling_loop(self):
+        # init text embeddings
+        text_ept = self.tokenizer('', **self.tokenizer_kwargs)
+        self.emb_ept = self.text_encoder(text_ept.input_ids.to(self.device))[0]
+        self.emb_spl = torch.cat([self.emb_ept, self.emb_pos, self.emb_neg], dim=0)
+        with torch.autocast(device_type=self.device, dtype=torch.float16):
+            # use noisy latent as starting point
+            x = self.latents[self.scheduler.timesteps[0].item()]
+            # sampling loop
+            for t in tqdm(self.scheduler.timesteps):
+                # concat latents & register timestep
+                src_latent = self.latents[t.item()]
+                latents = torch.cat([src_latent, x, x])
+                register_time(self, t.item())
+                # apply U-Net for denoising
+                noise_pred = self.unet(latents, t, encoder_hidden_states=self.emb_spl).sample
+                # classifier-free guidance
+                _, noise_pred_pos, noise_pred_neg = noise_pred.chunk(3)
+                noise_pred = noise_pred_neg + self.omega * (noise_pred_pos - noise_pred_neg)
+                # denoise step
+                x = self.scheduler.step(noise_pred, t, x).prev_sample
+            # save output latent
+            self.output_latent = x
+
+    def run_ditail(self):
+        # init output dir & dump config
+        os.makedirs(self.output_dir, exist_ok=True)
+        # self.save_dir = get_save_dir(self.output_dir)
+        # os.makedirs(self.save_dir, exist_ok=True)
+        # with open(os.path.join(self.output_dir, 'config.yaml'), 'w') as f:
+        #     if isinstance(self.args, dict):
+        #         f.write(yaml.dump(self.args))
+        #     else:
+        #         f.write(yaml.dump(vars(self.args)))
+        # step 1: inversion stage
+        self.load_inv_model()
+        self.extract_latents()
+        # self.latent_to_image(
+        #     latent=self.noisy_latent,
+        #     save_path=os.path.join(self.save_dir, 'noise.png')
+        # )
+        # step 2: sampling stage
+        self.load_spl_model()
+        if not self.no_injection:
+            self.init_injection()
+        self.sampling_loop()
+        return self.latent_to_image(
+            latent=self.output_latent,
+            # save_path=os.path.join(self.save_dir, 'output.png')
+        )
+
+def main(args):
+    seed_everything(args.seed)
+    ditail = DitailDemo(args)
+    ditail.run_ditail()
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--seed', type=int, default=42)
+    parser.add_argument('--device', type=str, default='cuda')
+    parser.add_argument('--output_dir', type=str, default='./output_demo')
+    parser.add_argument('--inv_model', type=str, default='runwayml/stable-diffusion-v1-5',
+                        help='Pre-trained inversion model name or path (step 1)')
+    parser.add_argument('--spl_model', type=str, default='runwayml/stable-diffusion-v1-5',
+                        help='Pre-trained sampling model name or path (step 2)')
+    parser.add_argument('--inv_steps', type=int, default=50,
+                        help='Number of inversion steps (step 1)')
+    parser.add_argument('--spl_steps', type=int, default=50,
+                        help='Number of sampling steps (step 2)')
+    # parser.add_argument('--img_path', type=str, required=True,
+    #                     help='Path to the source image')
+    parser.add_argument('--pos_prompt', type=str, required=True,
+                        help='Positive prompt for inversion')
+    parser.add_argument('--neg_prompt', type=str, default='worst quality, blurry, low res, NSFW',
+                        help='Negative prompt for inversion')
+    parser.add_argument('--alpha', type=float, default=2.0,
+                        help='Positive prompt scaling factor')
+    parser.add_argument('--beta', type=float, default=1.0,
+                        help='Negative prompt scaling factor')
+    parser.add_argument('--omega', type=float, default=15,
+                        help='Classifier-free guidance factor')
+    parser.add_argument('--mask', type=str, default='none',
+                        help='Optional mask for regional injection')
+    parser.add_argument('--lora', type=str, default='none',
+                        help='Optional LoRA for the sampling stage')
+    parser.add_argument('--lora_dir', type=str, default='./lora',
+                        help='Optional LoRA storing directory')
+    parser.add_argument('--lora_scale', type=float, default=0.7,
+                        help='Optional LoRA scaling weight')
+    parser.add_argument('--no_injection', action="store_true",
+                        help='Do not use PnP injection')
+    args = parser.parse_args()
+    main(args)

ditail/src/ditail_utils.py

@@ -0,0 +1,121 @@
+# credits: https://github.com/MichalGeyer/pnp-diffusers/blob/main/pnp_utils.py
+
+import os
+import torch
+import random
+import numpy as np
+
+def seed_everything(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+
+# def get_save_dir(output_dir, img_path):
+#     os.makedirs(output_dir, exist_ok=True)
+#     file = os.path.basename(img_path)
+#     indices = [d for d in os.listdir(output_dir) if d.startswith(file)]
+#     return os.path.join(output_dir, f'{file}_{len(indices)}')
+
+def register_time(model, t):
+    conv_module = model.unet.up_blocks[1].resnets[1]
+    setattr(conv_module, 't', t)
+    down_res_dict = {0: [0, 1], 1: [0, 1], 2: [0, 1]}
+    up_res_dict = {1: [0, 1, 2], 2: [0, 1, 2], 3: [0, 1, 2]}
+    for res in up_res_dict:
+        for block in up_res_dict[res]:
+            module = model.unet.up_blocks[res].attentions[block].transformer_blocks[0].attn1
+            setattr(module, 't', t)
+    for res in down_res_dict:
+        for block in down_res_dict[res]:
+            module = model.unet.down_blocks[res].attentions[block].transformer_blocks[0].attn1
+            setattr(module, 't', t)
+    module = model.unet.mid_block.attentions[0].transformer_blocks[0].attn1
+    setattr(module, 't', t)
+
+def register_attn_inj(model, injection_schedule):
+    def sa_forward(self):
+        to_out = self.to_out
+        if type(to_out) is torch.nn.modules.container.ModuleList:
+            to_out = self.to_out[0]
+        else:
+            to_out = self.to_out
+        def forward(x, encoder_hidden_states=None, attention_mask=None):
+            batch_size, sequence_length, dim = x.shape
+            h = self.heads
+            is_cross = encoder_hidden_states is not None
+            encoder_hidden_states = encoder_hidden_states if is_cross else x
+            q = self.to_q(x)
+            k = self.to_k(encoder_hidden_states)
+            v = self.to_v(encoder_hidden_states)
+            if not is_cross and self.injection_schedule is not None and (
+                    self.t in self.injection_schedule or self.t == 1000):
+                source_batch_size = int(q.shape[0] // 3)
+                # inject pos chunk
+                q[source_batch_size:2 * source_batch_size] = q[:source_batch_size]
+                k[source_batch_size:2 * source_batch_size] = k[:source_batch_size]
+                # inject neg chunk
+                q[2 * source_batch_size:] = q[:source_batch_size]
+                k[2 * source_batch_size:] = k[:source_batch_size]
+            q = self.head_to_batch_dim(q)
+            k = self.head_to_batch_dim(k)
+            v = self.head_to_batch_dim(v)
+            sim = torch.einsum("b i d, b j d -> b i j", q, k) * self.scale
+            if attention_mask is not None:
+                attention_mask = attention_mask.reshape(batch_size, -1)
+                max_neg_value = -torch.finfo(sim.dtype).max
+                attention_mask = attention_mask[:, None, :].repeat(h, 1, 1)
+                sim.masked_fill_(~attention_mask, max_neg_value)
+            attn = sim.softmax(dim=-1)
+            out = torch.einsum("b i j, b j d -> b i d", attn, v)
+            out = self.batch_to_head_dim(out)
+            return to_out(out)
+        return forward
+    res_dict = {1: [1, 2], 2: [0, 1, 2], 3: [0, 1, 2]}
+    for res in res_dict:
+        for block in res_dict[res]:
+            module = model.unet.up_blocks[res].attentions[block].transformer_blocks[0].attn1
+            module.forward = sa_forward(module)
+            setattr(module, 'injection_schedule', injection_schedule)
+
+def register_conv_inj(model, injection_schedule):
+    def conv_forward(self):
+        def forward(input_tensor, temb, scale):
+            hidden_states = input_tensor
+            hidden_states = self.norm1(hidden_states)
+            hidden_states = self.nonlinearity(hidden_states)
+            if self.upsample is not None:
+                if hidden_states.shape[0] >= 64:
+                    input_tensor = input_tensor.contiguous()
+                    hidden_states = hidden_states.contiguous()
+                input_tensor = self.upsample(input_tensor, scale=scale)
+                hidden_states = self.upsample(hidden_states, scale=scale)
+            elif self.downsample is not None:
+                input_tensor = self.downsample(input_tensor, scale=scale)
+                hidden_states = self.downsample(hidden_states, scale=scale)
+            hidden_states = self.conv1(hidden_states, scale)
+            if temb is not None:
+                temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None, None]
+            if temb is not None and self.time_embedding_norm == "default":
+                hidden_states = hidden_states + temb
+            hidden_states = self.norm2(hidden_states)
+            if temb is not None and self.time_embedding_norm == "scale_shift":
+                scale, shift = torch.chunk(temb, 2, dim=1)
+                hidden_states = hidden_states * (1 + scale) + shift
+            hidden_states = self.nonlinearity(hidden_states)
+            hidden_states = self.dropout(hidden_states)
+            hidden_states = self.conv2(hidden_states, scale)
+            if self.injection_schedule is not None and (self.t in self.injection_schedule or self.t == 1000):
+                source_batch_size = int(hidden_states.shape[0] // 3)
+                # inject pos chunk
+                hidden_states[source_batch_size:2 * source_batch_size] = hidden_states[:source_batch_size]
+                # inject neg chunk
+                hidden_states[2 * source_batch_size:] = hidden_states[:source_batch_size]
+            if self.conv_shortcut is not None:
+                input_tensor = self.conv_shortcut(input_tensor, scale)
+            output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
+            return output_tensor
+        return forward
+    conv_module = model.unet.up_blocks[1].resnets[1]
+    conv_module.forward = conv_forward(conv_module)
+    setattr(conv_module, 'injection_schedule', injection_schedule)

requirements.txt

@@ -0,0 +1,10 @@
+gradio
+accelerate
+--extra-index-url https://download.pytorch.org/whl/cu121
+torch==2.1.0
+torchvision
+transformers==4.35.2
+diffusers==0.24.0
+xformers
+open_clip_torch
+clip-interrogator