support for different base models

StableDiffuser.py

@@ -1,4 +1,5 @@
 import argparse
+import traceback
 
 import torch
 from baukit import TraceDict
@@ -36,71 +37,68 @@ def default_parser():
 class StableDiffuser(torch.nn.Module):
 
     def __init__(self,
-                scheduler='LMS'
+                scheduler='LMS',
+                 repo_id_or_path="CompVis/stable-diffusion-v1-4",
         ):
 
         super().__init__()
 
         # Load the autoencoder model which will be used to decode the latents into image space.
         self.vae = AutoencoderKL.from_pretrained(
-            "CompVis/stable-diffusion-v1-4", subfolder="vae")
+            repo_id_or_path, subfolder="vae")
         
         # Load the tokenizer and text encoder to tokenize and encode the text.
         self.tokenizer = CLIPTokenizer.from_pretrained(
-            "openai/clip-vit-large-patch14")
+            repo_id_or_path, subfolder="tokenizer")
         self.text_encoder = CLIPTextModel.from_pretrained(
-            "openai/clip-vit-large-patch14")
+            repo_id_or_path, subfolder="text_encoder")
         
         # The UNet model for generating the latents.
         self.unet = UNet2DConditionModel.from_pretrained(
-            "CompVis/stable-diffusion-v1-4", subfolder="unet")
-        
-        self.feature_extractor = CLIPFeatureExtractor.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="feature_extractor")
-        self.safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="safety_checker")
+            repo_id_or_path, subfolder="unet")
+
+        try:
+            self.feature_extractor = CLIPFeatureExtractor.from_pretrained(repo_id_or_path, subfolder="feature_extractor")
+            self.safety_checker = StableDiffusionSafetyChecker.from_pretrained(repo_id_or_path, subfolder="safety_checker")
+        except Exception as error:
+            print(f"caught exception {error} making feature extractor / safety checker")
+            self.feature_extractor = None
+            self.safety_checker = None
 
         if scheduler == 'LMS':
             self.scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
         elif scheduler == 'DDIM':
-            self.scheduler = DDIMScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler")
+            self.scheduler = DDIMScheduler.from_pretrained(repo_id_or_path, subfolder="scheduler")
         elif scheduler == 'DDPM':
-            self.scheduler = DDPMScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler")    
+            self.scheduler = DDPMScheduler.from_pretrained(repo_id_or_path, subfolder="scheduler")
 
         self.eval()
 
     def get_noise(self, batch_size, img_size, generator=None):
-
         param = list(self.parameters())[0]
-
         return torch.randn(
             (batch_size, self.unet.in_channels, img_size // 8, img_size // 8),
             generator=generator).type(param.dtype).to(param.device)
 
     def add_noise(self, latents, noise, step):
-
         return self.scheduler.add_noise(latents, noise, torch.tensor([self.scheduler.timesteps[step]]))
 
     def text_tokenize(self, prompts):
-
         return self.tokenizer(prompts, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt")
 
     def text_detokenize(self, tokens):
-
         return [self.tokenizer.decode(token) for token in tokens if token != self.tokenizer.vocab_size - 1]
 
     def text_encode(self, tokens):
-
         return self.text_encoder(tokens.input_ids.to(self.unet.device))[0]
 
     def decode(self, latents):
-
         return self.vae.decode(1 / self.vae.config.scaling_factor * latents).sample
 
     def encode(self, tensors):
-
         return self.vae.encode(tensors).latent_dist.mode() * 0.18215
 
     def to_image(self, image):
-
         image = (image / 2 + 0.5).clamp(0, 1)
         image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
         images = (image * 255).round().astype("uint8")
@@ -112,25 +110,16 @@ class StableDiffuser(torch.nn.Module):
         self.scheduler.set_timesteps(n_steps, device=self.unet.device)
 
     def get_initial_latents(self, n_imgs, img_size, n_prompts, generator=None):
-
         noise = self.get_noise(n_imgs, img_size, generator=generator).repeat(n_prompts, 1, 1, 1)
-
         latents = noise * self.scheduler.init_noise_sigma
-
         return latents
 
-    def get_text_embeddings(self, prompts, n_imgs):
-
+    def get_text_embeddings(self, prompts, negative_prompts, n_imgs):
         text_tokens = self.text_tokenize(prompts)
-
         text_embeddings = self.text_encode(text_tokens)
-
-        unconditional_tokens = self.text_tokenize([""] * len(prompts))
-
+        unconditional_tokens = self.text_tokenize(negative_prompts)
         unconditional_embeddings = self.text_encode(unconditional_tokens)
-
         text_embeddings = torch.cat([unconditional_embeddings, text_embeddings]).repeat_interleave(n_imgs, dim=0)
-
         return text_embeddings
 
     def predict_noise(self,
@@ -174,9 +163,7 @@ class StableDiffuser(torch.nn.Module):
         trace = None
 
         for iteration in tqdm(range(start_iteration, end_iteration), disable=not show_progress):
-
             if trace_args:
-
                 trace = TraceDict(self, **trace_args)
 
             noise_pred = self.predict_noise(
@@ -189,17 +176,13 @@ class StableDiffuser(torch.nn.Module):
             output = self.scheduler.step(noise_pred, self.scheduler.timesteps[iteration], latents)
 
             if trace_args:
-
                 trace.close()
-
                 trace_steps.append(trace)
 
             latents = output.prev_sample
 
             if return_steps or iteration == end_iteration - 1:
-
                 output = output.pred_original_sample if pred_x0 else latents
-
                 if return_steps:
                     latents_steps.append(output.cpu())
                 else:
@@ -210,6 +193,7 @@ class StableDiffuser(torch.nn.Module):
     @torch.no_grad()
     def __call__(self,
                  prompts,
+                 negative_prompts,
                  img_size=512,
                  n_steps=50,
                  n_imgs=1,
@@ -221,17 +205,12 @@ class StableDiffuser(torch.nn.Module):
         assert 0 <= n_steps <= 1000
 
         if not isinstance(prompts, list):
-
             prompts = [prompts]
 
         self.set_scheduler_timesteps(n_steps)
-
         latents = self.get_initial_latents(n_imgs, img_size, len(prompts), generator=generator)
-
-        text_embeddings = self.get_text_embeddings(prompts,n_imgs=n_imgs)
-
+        text_embeddings = self.get_text_embeddings(prompts,negative_prompts,n_imgs=n_imgs)
         end_iteration = end_iteration or n_steps
-
         latents_steps, trace_steps = self.diffusion(
             latents,
             text_embeddings,
@@ -242,19 +221,18 @@ class StableDiffuser(torch.nn.Module):
         latents_steps = [self.decode(latents.to(self.unet.device)) for latents in latents_steps]
         images_steps = [self.to_image(latents) for latents in latents_steps]
 
-        for i in range(len(images_steps)):
-            self.safety_checker = self.safety_checker.float()
-            safety_checker_input = self.feature_extractor(images_steps[i], return_tensors="pt").to(latents_steps[0].device)
-            image, has_nsfw_concept = self.safety_checker(
-                images=latents_steps[i].float().cpu().numpy(), clip_input=safety_checker_input.pixel_values.float()
-            )
-
-            images_steps[i][0] = self.to_image(torch.from_numpy(image))[0]
+        if self.safety_checker is not None:
+            for i in range(len(images_steps)):
+                self.safety_checker = self.safety_checker.float()
+                safety_checker_input = self.feature_extractor(images_steps[i], return_tensors="pt").to(latents_steps[0].device)
+                image, has_nsfw_concept = self.safety_checker(
+                    images=latents_steps[i].float().cpu().numpy(), clip_input=safety_checker_input.pixel_values.float()
+                )
+                images_steps[i][0] = self.to_image(torch.from_numpy(image))[0]
 
         images_steps = list(zip(*images_steps))
 
         if trace_steps:
-
             return images_steps, trace_steps
 
         return images_steps
@@ -263,7 +241,6 @@ class StableDiffuser(torch.nn.Module):
 if __name__ == '__main__':
 
     parser = default_parser()
-
     args = parser.parse_args()
 
     diffuser = StableDiffuser(seed=args.seed, scheduler='DDIM').to(torch.device(args.device)).half()

app.py

@@ -1,20 +1,27 @@
 import gradio as gr
 import torch
+import os
 from finetuning import FineTunedModel
 from StableDiffuser import StableDiffuser
 from train import train
     
 import os
-model_map = {'Van Gogh' : 'models/vangogh.pt', 
+model_map = {'Van Gogh': 'models/vangogh.pt',
              'Pablo Picasso': 'models/pablopicasso.pt',
-             'Car' : 'models/car.pt',
+             'Car': 'models/car.pt',
              'Garbage Truck': 'models/garbagetruck.pt',
              'French Horn': 'models/frenchhorn.pt',
-             'Kilian Eng' : 'models/kilianeng.pt',
-             'Thomas Kinkade' : 'models/thomaskinkade.pt',
-             'Tyler Edlin' : 'models/tyleredlin.pt',
+             'Kilian Eng': 'models/kilianeng.pt',
+             'Thomas Kinkade': 'models/thomaskinkade.pt',
+             'Tyler Edlin': 'models/tyleredlin.pt',
              'Kelly McKernan': 'models/kellymckernan.pt',
              'Rembrandt': 'models/rembrandt.pt' }
+for model_file in os.listdir('models'):
+    path = 'models/' + model_file
+    if any([existing_path == path for existing_path in model_map.values()]):
+        continue
+    model_map[model_file] = path
+
 
 ORIGINAL_SPACE_ID = 'baulab/Erasing-Concepts-In-Diffusion'
 SPACE_ID = os.getenv('SPACE_ID')
@@ -31,8 +38,6 @@ class Demo:
         self.training = False
         self.generating = False
 
-        self.diffuser = StableDiffuser(scheduler='DDIM').to('cuda').eval().half()
-
         with gr.Blocks() as demo:
             self.layout()
             demo.queue(concurrency_count=5).launch()
@@ -64,6 +69,9 @@ class Demo:
                             label="Prompt",
                             info="Prompt to generate"
                         )
+                        self.negative_prompt_input_infr = gr.Text(
+                            label="Negative prompt"
+                        )
 
                         with gr.Row():
 
@@ -78,6 +86,19 @@ class Demo:
                                 label="Seed",
                                 value=42
                             )
+                            self.img_size_infr = gr.Slider(
+                                label="Image size",
+                                minimum=256,
+                                maximum=1024,
+                                value=512,
+                                step=64
+                            )
+
+                        self.base_repo_id_or_path_input_infr = gr.Text(
+                            label="Base model",
+                            value="CompVis/stable-diffusion-v1-4",
+                            info="Path or huggingface repo id of the base model that this edit was done against"
+                        )
 
                     with gr.Column(scale=2):
 
@@ -108,6 +129,21 @@ class Demo:
 
                     with gr.Column(scale=3):
 
+                        self.train_model_input = gr.Text(
+                            label="Model to Edit",
+                            value="CompVis/stable-diffusion-v1-4",
+                            info="Path or huggingface repo id of the model to edit"
+                        )
+
+                        self.train_img_size_input = gr.Slider(
+                            value=512,
+                            step=64,
+                            minimum=256,
+                            maximum=1024,
+                            label="Image Size",
+                            info="Image size for training, should match the model's native image size"
+                        )
+
                         self.prompt_input = gr.Text(
                             placeholder="Enter prompt...",
                             label="Prompt to Erase",
@@ -156,8 +192,11 @@ class Demo:
 
         self.infr_button.click(self.inference, inputs = [
             self.prompt_input_infr,
+            self.negative_prompt_input_infr,
             self.seed_infr,
-            self.model_dropdown
+            self.img_size_infr,
+            self.model_dropdown,
+            self.base_repo_id_or_path_input_infr
             ],
             outputs=[
                 self.image_new,
@@ -165,6 +204,8 @@ class Demo:
             ]
         )
         self.train_button.click(self.train, inputs = [
+            self.train_model_input,
+            self.train_img_size_input,
             self.prompt_input,
             self.train_method_input, 
             self.neg_guidance_input,
@@ -174,7 +215,7 @@ class Demo:
         outputs=[self.train_button,  self.train_status, self.download, self.model_dropdown]
         )
 
-    def train(self, prompt, train_method, neg_guidance, iterations, lr, pbar = gr.Progress(track_tqdm=True)):
+    def train(self, repo_id_or_path, img_size, prompt, train_method, neg_guidance, iterations, lr, pbar = gr.Progress(track_tqdm=True)):
 
         if self.training:
             return [gr.update(interactive=True, value='Train'), gr.update(value='Someone else is training... Try again soon'), None, gr.update()]
@@ -200,7 +241,7 @@ class Demo:
 
         self.training = True
 
-        train(prompt, modules, frozen, iterations, neg_guidance, lr, save_path)
+        train(repo_id_or_path, img_size, prompt, modules, frozen, iterations, neg_guidance, lr, save_path)
 
         self.training = False
 
@@ -211,22 +252,21 @@ class Demo:
         return [gr.update(interactive=True, value='Train'), gr.update(value='Done Training! \n Try your custom model in the "Test" tab'), save_path, gr.Dropdown.update(choices=list(model_map.keys()), value='Custom')]
 
 
-    def inference(self, prompt, seed, model_name, pbar = gr.Progress(track_tqdm=True)):
+    def inference(self, prompt, negative_prompt, seed, img_size, model_name, base_repo_id_or_path, pbar = gr.Progress(track_tqdm=True)):
         
         seed = seed or 42
-
         generator = torch.manual_seed(seed)
-
         model_path = model_map[model_name]
-        
         checkpoint = torch.load(model_path)
 
+        self.diffuser = StableDiffuser(scheduler='DDIM', repo_id_or_path=base_repo_id_or_path).to('cuda').eval().half()
         finetuner = FineTunedModel.from_checkpoint(self.diffuser, checkpoint).eval().half()
-
         torch.cuda.empty_cache()
 
         images = self.diffuser(
             prompt,
+            negative_prompt,
+            img_size=img_size,
             n_steps=50,
             generator=generator
         )
@@ -242,6 +282,7 @@ class Demo:
 
             images = self.diffuser(
                 prompt,
+                negative_prompt,
                 n_steps=50,
                 generator=generator
             )

train.py

@@ -3,11 +3,11 @@ from finetuning import FineTunedModel
 import torch
 from tqdm import tqdm
 
-def train(prompt, modules, freeze_modules, iterations, negative_guidance, lr, save_path):
+def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations, negative_guidance, lr, save_path):
   
     nsteps = 50
 
-    diffuser = StableDiffuser(scheduler='DDIM').to('cuda')
+    diffuser = StableDiffuser(scheduler='DDIM', repo_id_or_path=repo_id_or_path).to('cuda')
     diffuser.train()
 
     finetuner = FineTunedModel(diffuser, modules, frozen_modules=freeze_modules)
@@ -28,17 +28,16 @@ def train(prompt, modules, freeze_modules, iterations, negative_guidance, lr, sa
 
     torch.cuda.empty_cache()
 
+    print(f"using img_size of {img_size}")
+
     for i in pbar:
-        
         with torch.no_grad():
-
             diffuser.set_scheduler_timesteps(nsteps)
-
             optimizer.zero_grad()
 
             iteration = torch.randint(1, nsteps - 1, (1,)).item()
 
-            latents = diffuser.get_initial_latents(1, 512, 1)
+            latents = diffuser.get_initial_latents(1, img_size, 1)
 
             with finetuner:
 
@@ -80,6 +79,8 @@ if __name__ == '__main__':
 
     parser = argparse.ArgumentParser()
 
+    parser.add_argument("--repo_id_or_path", required=True)
+    parser.add_argument("--img_size", type=int, required=False, default=512)
     parser.add_argument('--prompt', required=True)
     parser.add_argument('--modules', required=True)
     parser.add_argument('--freeze_modules', nargs='+', required=True)