allow different resolutions for w/h

StableDiffuser.py

@@ -1,17 +1,13 @@
 import argparse
-import traceback
 
 import torch
 from baukit import TraceDict
-from diffusers import AutoencoderKL, UNet2DConditionModel
+from diffusers import StableDiffusionPipeline
 from PIL import Image
 from tqdm.auto import tqdm
-from transformers import CLIPTextModel, CLIPTokenizer, CLIPFeatureExtractor
-from diffusers.schedulers import EulerAncestralDiscreteScheduler
 from diffusers.schedulers.scheduling_ddim import DDIMScheduler
 from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
 from diffusers.schedulers.scheduling_lms_discrete import LMSDiscreteScheduler
-from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
 import util
 
 
@@ -39,31 +35,17 @@ class StableDiffuser(torch.nn.Module):
     def __init__(self,
                 scheduler='LMS',
                  repo_id_or_path="CompVis/stable-diffusion-v1-4",
+                 variant='fp16'
         ):
 
         super().__init__()
 
-        # Load the autoencoder model which will be used to decode the latents into image space.
-        self.vae = AutoencoderKL.from_pretrained(
-            repo_id_or_path, subfolder="vae")
-        
-        # Load the tokenizer and text encoder to tokenize and encode the text.
-        self.tokenizer = CLIPTokenizer.from_pretrained(
-            repo_id_or_path, subfolder="tokenizer")
-        self.text_encoder = CLIPTextModel.from_pretrained(
-            repo_id_or_path, subfolder="text_encoder")
-        
-        # The UNet model for generating the latents.
-        self.unet = UNet2DConditionModel.from_pretrained(
-            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
+        self.pipeline = StableDiffusionPipeline.from_pretrained(repo_id_or_path, variant=variant)
+
+        self.vae = self.pipeline.vae
+        self.unet = self.pipeline.unet
+        self.tokenizer = self.pipeline.tokenizer
+        self.text_encoder = self.pipeline.text_encoder
 
         if scheduler == 'LMS':
             self.scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
@@ -74,10 +56,14 @@ class StableDiffuser(torch.nn.Module):
 
         self.eval()
 
-    def get_noise(self, batch_size, img_size, generator=None):
+    @property
+    def safety_checker(self):
+        return self.pipeline.safety_checker
+
+    def get_noise(self, batch_size, width, height, generator=None):
         param = list(self.parameters())[0]
         return torch.randn(
-            (batch_size, self.unet.in_channels, img_size // 8, img_size // 8),
+            (batch_size, self.unet.in_channels, width // 8, height // 8),
             generator=generator).type(param.dtype).to(param.device)
 
     def add_noise(self, latents, noise, step):
@@ -109,8 +95,8 @@ class StableDiffuser(torch.nn.Module):
     def set_scheduler_timesteps(self, n_steps):
         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)
+    def get_initial_latents(self, n_imgs, width, height, n_prompts, generator=None):
+        noise = self.get_noise(n_imgs, width, height, generator=generator).repeat(n_prompts, 1, 1, 1)
         latents = noise * self.scheduler.init_noise_sigma
         return latents
 
@@ -196,7 +182,8 @@ class StableDiffuser(torch.nn.Module):
     def __call__(self,
                  prompts,
                  negative_prompts,
-                 img_size=512,
+                 width=512,
+                 height=512,
                  n_steps=50,
                  n_imgs=1,
                  end_iteration=None,
@@ -210,7 +197,7 @@ class StableDiffuser(torch.nn.Module):
             prompts = [prompts]
 
         self.set_scheduler_timesteps(n_steps)
-        latents = self.get_initial_latents(n_imgs, img_size, len(prompts), generator=generator)
+        latents = self.get_initial_latents(n_imgs, width, height, len(prompts), generator=generator)
         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(
@@ -239,13 +226,16 @@ class StableDiffuser(torch.nn.Module):
 
         return images_steps
 
+    def save_pretrained(self, path, **kwargs):
+        self.pipeline.save_pretrained(path, **kwargs)
+
 
 if __name__ == '__main__':
 
     parser = default_parser()
     args = parser.parse_args()
 
-    diffuser = StableDiffuser(seed=args.seed, scheduler='DDIM').to(torch.device(args.device)).half()
+    diffuser = StableDiffuser(scheduler='DDIM').to(torch.device(args.device)).half()
 
     images = diffuser(args.prompts,
                       n_steps=args.nsteps,

app.py

@@ -86,8 +86,16 @@ class Demo:
                                 label="Seed",
                                 value=42
                             )
-                            self.img_size_infr = gr.Slider(
-                                label="Image size",
+                            self.img_width_infr = gr.Slider(
+                                label="Image width",
+                                minimum=256,
+                                maximum=1024,
+                                value=512,
+                                step=64
+                            )
+
+                            self.img_height_infr = gr.Slider(
+                                label="Image height",
                                 minimum=256,
                                 maximum=1024,
                                 value=512,
@@ -190,11 +198,51 @@ class Demo:
 
                         self.download = gr.Files()
 
+            with gr.Tab("Export") as export_column:
+
+                with gr.Row():
+
+                    self.explain_train= gr.Markdown(interactive=False,
+                                      value='Export a model to Diffusers format. Please enter the base model and select the editing weights.')
+
+                with gr.Row():
+
+                    with gr.Column(scale=3):
+
+                        self.base_repo_id_or_path_input_export = 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"
+                        )
+
+                        self.model_dropdown_export = gr.Dropdown(
+                            label="ESD Model",
+                            choices=list(model_map.keys()),
+                            value='Van Gogh',
+                            interactive=True
+                        )
+
+                        self.save_path_input_export = gr.Text(
+                            label="Output path",
+                            placeholder="./exported_models/model_name",
+                            info="Path to export the model to. A diffusers folder will be written to this location."
+                        )
+
+                        self.save_half_export = gr.Checkbox(
+                            label="Save as fp16"
+                        )
+
+                    with gr.Column(scale=1):
+                        self.export_button = gr.Button(
+                            value="Export",
+                        )
+
         self.infr_button.click(self.inference, inputs = [
             self.prompt_input_infr,
             self.negative_prompt_input_infr,
             self.seed_infr,
-            self.img_size_infr,
+            self.img_width_infr,
+            self.img_height_infr,
             self.model_dropdown,
             self.base_repo_id_or_path_input_infr
             ],
@@ -214,6 +262,14 @@ class Demo:
         ],
         outputs=[self.train_button,  self.train_status, self.download, self.model_dropdown]
         )
+        self.export_button.click(self.export, inputs = [
+            self.model_dropdown_export,
+            self.base_repo_id_or_path_input_export,
+            self.save_path_input_export,
+            self.save_half_export
+        ],
+        outputs=[self.export_button]
+        )
 
     def train(self, repo_id_or_path, img_size, prompt, train_method, neg_guidance, iterations, lr, pbar = gr.Progress(track_tqdm=True)):
 
@@ -251,42 +307,50 @@ 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 export(self, model_name, base_repo_id_or_path, save_path, save_half):
+        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()
+        finetuner = FineTunedModel.from_checkpoint(self.diffuser, checkpoint).eval()
+        with finetuner:
+            if save_half:
+                self.diffuser = self.diffuser.half()
+                self.diffuser.pipeline.to(torch.float16, torch_device=self.diffuser.device)
+            self.diffuser.save_pretrained(save_path)
+
 
-    def inference(self, prompt, negative_prompt, seed, img_size, model_name, base_repo_id_or_path, pbar = gr.Progress(track_tqdm=True)):
+    def inference(self, prompt, negative_prompt, seed, width, height, 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()
 
+        generator = torch.manual_seed(seed)
+
+        torch.cuda.empty_cache()
         images = self.diffuser(
             prompt,
             negative_prompt,
-            img_size=img_size,
+            width=width,
+            height=height,
             n_steps=50,
             generator=generator
         )
-
-        
         orig_image = images[0][0]
 
         torch.cuda.empty_cache()
-
-        generator = torch.manual_seed(seed)
-
         with finetuner:
-
             images = self.diffuser(
                 prompt,
                 negative_prompt,
+                width=width,
+                height=height,
                 n_steps=50,
                 generator=generator
             )
-
         edited_image = images[0][0]
 
         del finetuner

finetuning.py

@@ -2,11 +2,12 @@ import copy
 import re
 import torch
 import util
+from StableDiffuser import StableDiffuser
 
 class FineTunedModel(torch.nn.Module):
 
     def __init__(self,
-                 model,
+                 model: StableDiffuser,
                  modules,
                  frozen_modules=[]
                  ):
@@ -24,11 +25,8 @@ class FineTunedModel(torch.nn.Module):
 
         for module_name, module in model.named_modules():
             for ft_module_regex in modules:
-
                 match = re.search(ft_module_regex, module_name)
-                
                 if match is not None:
-
                     ft_module = copy.deepcopy(module)
                     
                     self.orig_modules[module_name] = module
@@ -39,13 +37,10 @@ class FineTunedModel(torch.nn.Module):
                     print(f"=> Finetuning {module_name}")
        
                     for ft_module_name, module in ft_module.named_modules():
-
                         ft_module_name = f"{module_name}.{ft_module_name}"
-
                         for freeze_module_name in frozen_modules:
 
                             match = re.search(freeze_module_name, ft_module_name)
-
                             if match:
                                 print(f"=> Freezing {ft_module_name}")
                                 util.freeze(module)

train.py

@@ -36,7 +36,7 @@ def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations
             optimizer.zero_grad()
 
             iteration = torch.randint(1, nsteps - 1, (1,)).item()
-            latents = diffuser.get_initial_latents(1, img_size, 1)
+            latents = diffuser.get_initial_latents(1, width=img_size, height=img_size, n_prompts=1)
 
             with finetuner:
                 latents_steps, _ = diffuser.diffusion(