Duplicate from baulab/Erasing-Concepts-In-Diffusion

Co-authored-by: Jaden Fiotto-Kaufman <[email protected]>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+images/applications.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+__pycache__

README.md

@@ -0,0 +1,54 @@
+---
+title: Erasing Concepts from Diffusion Models
+emoji: 💡
+colorFrom: indigo
+colorTo: gray
+sdk: gradio
+sdk_version: 3.21.0
+app_file: app.py
+pinned: false
+license: mit
+duplicated_from: baulab/Erasing-Concepts-In-Diffusion
+---
+
+
+
+# Erasing Concepts from Diffusion Models
+
+ Project Website    [https://erasing.baulab.info](https://erasing.baulab.info) <br>
+ Arxiv Preprint     [https://arxiv.org/pdf/2303.07345.pdf](https://arxiv.org/pdf/2303.07345.pdf) <br>
+ Fine-tuned Weights [https://erasing.baulab.info/weights/esd_models/](https://erasing.baulab.info/weights/esd_models/) <br>
+<div align='center'>
+<img src = 'images/applications.png'>
+</div>
+
+Motivated by recent advancements in text-to-image diffusion, we study erasure of specific concepts from the model's weights. While Stable Diffusion has shown promise in producing explicit or realistic artwork, it has raised concerns regarding its potential for misuse. We propose a fine-tuning method that can erase a visual concept from a pre-trained diffusion model, given only the name of the style and using negative guidance as a teacher. We benchmark our method against previous approaches that remove sexually explicit content and demonstrate its effectiveness, performing on par with Safe Latent Diffusion and censored training.
+
+To evaluate artistic style removal, we conduct experiments erasing five modern artists from the network and conduct a user study to assess the human perception of the removed styles. Unlike previous methods, our approach can remove concepts from a diffusion model permanently rather than modifying the output at the inference time, so it cannot be circumvented even if a user has access to model weights
+
+Given only a short text description of an undesired visual concept and no additional data, our method fine-tunes model weights to erase the targeted concept. Our method can avoid NSFW content, stop imitation of a specific artist's style, or even erase a whole object class from model output, while preserving the model's behavior and capabilities on other topics.
+
+## Demo vs github
+
+This demo uses an updated implementation from the original Erasing codebase the publication is based from.
+
+## Running locally
+
+1.) Create an environment using the packages included in the requirements.txt file
+
+2.) Run `python app.py`
+
+3.) Open the application in browser at `http://127.0.0.1:7860/`
+
+4.) Train, evaluate, and save models using our method
+
+## Citing our work
+The preprint can be cited as follows
+```
+@article{gandikota2023erasing,
+  title={Erasing Concepts from Diffusion Models},
+  author={Rohit Gandikota and Joanna Materzy\'nska and Jaden Fiotto-Kaufman and David Bau},
+  journal={arXiv preprint arXiv:2303.07345},
+  year={2023}
+}
+```

StableDiffuser.py

@@ -0,0 +1,279 @@
+import argparse
+
+import torch
+from baukit import TraceDict
+from diffusers import AutoencoderKL, UNet2DConditionModel
+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
+
+
+def default_parser():
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('prompts', type=str, nargs='+')
+    parser.add_argument('outpath', type=str)
+
+    parser.add_argument('--images', type=str, nargs='+', default=None)
+    parser.add_argument('--nsteps', type=int, default=1000)
+    parser.add_argument('--nimgs', type=int, default=1)
+    parser.add_argument('--start_itr', type=int, default=0)
+    parser.add_argument('--return_steps', action='store_true', default=False)
+    parser.add_argument('--pred_x0', action='store_true', default=False)
+    parser.add_argument('--device', type=str, default='cuda:0')
+    parser.add_argument('--seed', type=int, default=42)
+
+    return parser
+
+
+class StableDiffuser(torch.nn.Module):
+
+    def __init__(self,
+                scheduler='LMS'
+        ):
+
+        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")
+        
+        # Load the tokenizer and text encoder to tokenize and encode the text.
+        self.tokenizer = CLIPTokenizer.from_pretrained(
+            "openai/clip-vit-large-patch14")
+        self.text_encoder = CLIPTextModel.from_pretrained(
+            "openai/clip-vit-large-patch14")
+        
+        # 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")
+
+        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")
+        elif scheduler == 'DDPM':
+            self.scheduler = DDPMScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", 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")
+        pil_images = [Image.fromarray(image) for image in images]
+
+        return pil_images
+
+    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)
+
+        latents = noise * self.scheduler.init_noise_sigma
+
+        return latents
+
+    def get_text_embeddings(self, prompts, n_imgs):
+
+        text_tokens = self.text_tokenize(prompts)
+
+        text_embeddings = self.text_encode(text_tokens)
+
+        unconditional_tokens = self.text_tokenize([""] * len(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,
+             iteration,
+             latents,
+             text_embeddings,
+             guidance_scale=7.5
+             ):
+
+        # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
+        latents = torch.cat([latents] * 2)
+        latents = self.scheduler.scale_model_input(
+            latents, self.scheduler.timesteps[iteration])
+
+        # predict the noise residual
+        noise_prediction = self.unet(
+            latents, self.scheduler.timesteps[iteration], encoder_hidden_states=text_embeddings).sample
+
+        # perform guidance
+        noise_prediction_uncond, noise_prediction_text = noise_prediction.chunk(2)
+        noise_prediction = noise_prediction_uncond + guidance_scale * \
+            (noise_prediction_text - noise_prediction_uncond)
+
+        return noise_prediction
+
+    @torch.no_grad()
+    def diffusion(self,
+                  latents,
+                  text_embeddings,
+                  end_iteration=1000,
+                  start_iteration=0,
+                  return_steps=False,
+                  pred_x0=False,
+                  trace_args=None,                  
+                  show_progress=True,
+                  **kwargs):
+
+        latents_steps = []
+        trace_steps = []
+
+        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(
+                iteration, 
+                latents, 
+                text_embeddings,
+                **kwargs)
+
+            # compute the previous noisy sample x_t -> x_t-1
+            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:
+                    latents_steps.append(output)
+
+        return latents_steps, trace_steps
+
+    @torch.no_grad()
+    def __call__(self,
+                 prompts,
+                 img_size=512,
+                 n_steps=50,
+                 n_imgs=1,
+                 end_iteration=None,
+                 generator=None,
+                 **kwargs
+                 ):
+
+        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)
+
+        end_iteration = end_iteration or n_steps
+
+        latents_steps, trace_steps = self.diffusion(
+            latents,
+            text_embeddings,
+            end_iteration=end_iteration,
+            **kwargs
+        )
+
+        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]
+
+        images_steps = list(zip(*images_steps))
+
+        if trace_steps:
+
+            return images_steps, trace_steps
+
+        return images_steps
+
+
+if __name__ == '__main__':
+
+    parser = default_parser()
+
+    args = parser.parse_args()
+
+    diffuser = StableDiffuser(seed=args.seed, scheduler='DDIM').to(torch.device(args.device)).half()
+
+    images = diffuser(args.prompts,
+                      n_steps=args.nsteps,
+                      n_imgs=args.nimgs,
+                      start_iteration=args.start_itr,
+                      return_steps=args.return_steps,
+                      pred_x0=args.pred_x0
+                      )
+
+    util.image_grid(images, args.outpath)

app.py

@@ -0,0 +1,258 @@
+import gradio as gr
+import torch
+from finetuning import FineTunedModel
+from StableDiffuser import StableDiffuser
+from train import train
+    
+import os
+model_map = {'Van Gogh' : 'models/vangogh.pt', 
+             'Pablo Picasso': 'models/pablopicasso.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',
+             'Kelly McKernan': 'models/kellymckernan.pt',
+             'Rembrandt': 'models/rembrandt.pt' }
+
+ORIGINAL_SPACE_ID = 'baulab/Erasing-Concepts-In-Diffusion'
+SPACE_ID = os.getenv('SPACE_ID')
+
+SHARED_UI_WARNING = f'''## Attention - Training using the ESD-u method does not work in this shared UI. You can either duplicate and use it with a gpu with at least 40GB, or clone this repository to run on your own machine.
+<center><a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true"><img style="margin-top:0;margin-bottom:0" src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></center>
+'''
+
+
+class Demo:
+
+    def __init__(self) -> None:
+
+        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()
+
+
+    def layout(self):
+
+        with gr.Row():
+
+            if SPACE_ID == ORIGINAL_SPACE_ID:
+
+                self.warning = gr.Markdown(SHARED_UI_WARNING)
+          
+        with gr.Row():
+                
+            with gr.Tab("Test") as inference_column:
+
+                with gr.Row():
+
+                    self.explain_infr = gr.Markdown(interactive=False, 
+                                      value='This is a demo of [Erasing Concepts from Stable Diffusion](https://erasing.baulab.info/).  To try out a model where a concept has been erased, select a model and enter any prompt.  For example, if you select the model "Van Gogh" you can generate images for the prompt "A portrait in the style of Van Gogh" and compare the erased and unerased models.  We have also provided several other pre-fine-tuned models with artistic styles and objects erased (Check out the "ESD Model" drop-down). You can also train and run your own custom models. Check out the "train" section for custom erasure of concepts.')
+
+                with gr.Row():
+
+                    with gr.Column(scale=1):
+
+                        self.prompt_input_infr = gr.Text(
+                            placeholder="Enter prompt...",
+                            label="Prompt",
+                            info="Prompt to generate"
+                        )
+
+                        with gr.Row():
+
+                            self.model_dropdown = gr.Dropdown(
+                                label="ESD Model",
+                                choices= list(model_map.keys()),
+                                value='Van Gogh',
+                                interactive=True
+                            )
+
+                            self.seed_infr = gr.Number(
+                                label="Seed",
+                                value=42
+                            )
+
+                    with gr.Column(scale=2):
+
+                        self.infr_button = gr.Button(
+                            value="Generate",
+                            interactive=True
+                        )
+
+                        with gr.Row():
+
+                            self.image_new = gr.Image(
+                                label="ESD",
+                                interactive=False
+                            )
+                            self.image_orig = gr.Image(
+                                label="SD",
+                                interactive=False
+                            )
+
+            with gr.Tab("Train") as training_column:
+
+                with gr.Row():
+
+                    self.explain_train= gr.Markdown(interactive=False, 
+                                      value='In this part you can erase any concept from Stable Diffusion.   Enter a prompt for the concept or style you want to erase, and select ESD-x if you want to focus erasure on prompts that mention the concept explicitly. [NOTE: ESD-u is currently unavailable in this space. But you can duplicate the space and run it on GPU with VRAM >40GB for enabling ESD-u]. With default settings, it takes about 15 minutes to fine-tune the model; then you can try inference above or download the weights.  The training code used here is slightly different than the code tested in the original paper.  Code and details are at [github link](https://github.com/rohitgandikota/erasing).')
+
+                with gr.Row():
+
+                    with gr.Column(scale=3):
+
+                        self.prompt_input = gr.Text(
+                            placeholder="Enter prompt...",
+                            label="Prompt to Erase",
+                            info="Prompt corresponding to concept to erase"
+                        )
+
+                        choices = ['ESD-x']
+                        if torch.cuda.get_device_properties(0).total_memory * 1e-9 >= 40:
+                            choices.append('ESD-u')
+                    
+                        self.train_method_input = gr.Dropdown(
+                            choices=choices,
+                            value='ESD-x',
+                            label='Train Method',
+                            info='Method of training'
+                        )
+
+                        self.neg_guidance_input = gr.Number(
+                            value=1,
+                            label="Negative Guidance",
+                            info='Guidance of negative training used to train'
+                        )
+
+                        self.iterations_input = gr.Number(
+                            value=150,
+                            precision=0,
+                            label="Iterations",
+                            info='iterations used to train'
+                        )
+
+                        self.lr_input = gr.Number(
+                            value=1e-5,
+                            label="Learning Rate",
+                            info='Learning rate used to train'
+                        )
+
+                    with gr.Column(scale=1):
+
+                        self.train_status = gr.Button(value='', variant='primary', label='Status', interactive=False)
+
+                        self.train_button = gr.Button(
+                            value="Train",
+                        )
+
+                        self.download = gr.Files()
+
+        self.infr_button.click(self.inference, inputs = [
+            self.prompt_input_infr,
+            self.seed_infr,
+            self.model_dropdown
+            ],
+            outputs=[
+                self.image_new,
+                self.image_orig
+            ]
+        )
+        self.train_button.click(self.train, inputs = [
+            self.prompt_input,
+            self.train_method_input, 
+            self.neg_guidance_input,
+            self.iterations_input,
+            self.lr_input
+        ],
+        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)):
+
+        if self.training:
+            return [gr.update(interactive=True, value='Train'), gr.update(value='Someone else is training... Try again soon'), None, gr.update()]
+
+        if train_method == 'ESD-x':
+
+            modules = ".*attn2$"
+            frozen = []
+
+        elif train_method == 'ESD-u':
+
+            modules = "unet$"
+            frozen = [".*attn2$", "unet.time_embedding$", "unet.conv_out$"]   
+
+        elif train_method == 'ESD-self':
+
+            modules = ".*attn1$"
+            frozen = []
+
+        randn = torch.randint(1, 10000000, (1,)).item()
+
+        save_path = f"models/{randn}_{prompt.lower().replace(' ', '')}.pt"
+
+        self.training = True
+
+        train(prompt, modules, frozen, iterations, neg_guidance, lr, save_path)
+
+        self.training = False
+
+        torch.cuda.empty_cache()
+
+        model_map['Custom'] = save_path
+
+        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)):
+        
+        seed = seed or 42
+
+        generator = torch.manual_seed(seed)
+
+        model_path = model_map[model_name]
+        
+        checkpoint = torch.load(model_path)
+
+        finetuner = FineTunedModel.from_checkpoint(self.diffuser, checkpoint).eval().half()
+
+        torch.cuda.empty_cache()
+
+        images = self.diffuser(
+            prompt,
+            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,
+                n_steps=50,
+                generator=generator
+            )
+
+        edited_image = images[0][0]
+
+        del finetuner
+        torch.cuda.empty_cache()
+
+        return edited_image, orig_image
+
+
+demo = Demo()
+

finetuning.py

@@ -0,0 +1,101 @@
+import copy
+import re
+import torch
+import util
+
+class FineTunedModel(torch.nn.Module):
+
+    def __init__(self,
+                 model,
+                 modules,
+                 frozen_modules=[]
+                 ):
+
+        super().__init__()
+
+        if isinstance(modules, str):
+            modules = [modules]
+
+        self.model = model
+        self.ft_modules = {}
+        self.orig_modules = {}
+
+        util.freeze(self.model)
+
+        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
+                    self.ft_modules[module_name] = ft_module
+
+                    util.unfreeze(ft_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)
+
+        self.ft_modules_list = torch.nn.ModuleList(self.ft_modules.values())
+        self.orig_modules_list = torch.nn.ModuleList(self.orig_modules.values())
+
+
+    @classmethod
+    def from_checkpoint(cls, model, checkpoint, frozen_modules=[]):
+
+        if isinstance(checkpoint, str):
+            checkpoint = torch.load(checkpoint)
+
+        modules = [f"{key}$" for key in list(checkpoint.keys())]
+
+        ftm = FineTunedModel(model, modules, frozen_modules=frozen_modules)
+        ftm.load_state_dict(checkpoint)
+
+        return ftm
+
+        
+    def __enter__(self):
+
+        for key, ft_module in self.ft_modules.items():
+            util.set_module(self.model, key, ft_module)
+
+    def __exit__(self, exc_type, exc_value, tb):
+
+        for key, module in self.orig_modules.items():
+            util.set_module(self.model, key, module)
+
+    def parameters(self):
+
+        parameters = []
+
+        for ft_module in self.ft_modules.values():
+
+            parameters.extend(list(ft_module.parameters()))
+
+        return parameters
+
+    def state_dict(self):
+
+        state_dict = {key: module.state_dict() for key, module in self.ft_modules.items()}
+
+        return state_dict
+
+    def load_state_dict(self, state_dict):
+
+        for key, sd in state_dict.items():
+            
+            self.ft_modules[key].load_state_dict(sd)

models/car.pt

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

models/frenchhorn.pt

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

models/garbagetruck.pt

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

models/kellymckernan.pt

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

models/kilianeng.pt

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

models/pablopicasso.pt

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

models/rembrandt.pt

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

models/thomaskinkade.pt

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

models/tyleredlin.pt

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

models/vangogh.pt

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

requirements.txt

@@ -0,0 +1,8 @@
+gradio
+torch==1.13.1 --index-url https://download.pytorch.org/whl/cu118
+torchvision==0.14.1 --index-url https://download.pytorch.org/whl/cu118
+diffusers
+transformers
+accelerate
+scipy
+git+https://github.com/davidbau/baukit.git

train.py

@@ -0,0 +1,91 @@
+from StableDiffuser import StableDiffuser
+from finetuning import FineTunedModel
+import torch
+from tqdm import tqdm
+
+def train(prompt, modules, freeze_modules, iterations, negative_guidance, lr, save_path):
+  
+    nsteps = 50
+
+    diffuser = StableDiffuser(scheduler='DDIM').to('cuda')
+    diffuser.train()
+
+    finetuner = FineTunedModel(diffuser, modules, frozen_modules=freeze_modules)
+
+    optimizer = torch.optim.Adam(finetuner.parameters(), lr=lr)
+    criteria = torch.nn.MSELoss()
+
+    pbar = tqdm(range(iterations))
+
+    with torch.no_grad():
+
+        neutral_text_embeddings = diffuser.get_text_embeddings([''],n_imgs=1)
+        positive_text_embeddings = diffuser.get_text_embeddings([prompt],n_imgs=1)
+
+    del diffuser.vae
+    del diffuser.text_encoder
+    del diffuser.tokenizer
+
+    torch.cuda.empty_cache()
+
+    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)
+
+            with finetuner:
+
+                latents_steps, _ = diffuser.diffusion(
+                    latents,
+                    positive_text_embeddings,
+                    start_iteration=0,
+                    end_iteration=iteration,
+                    guidance_scale=3, 
+                    show_progress=False
+                )
+
+            diffuser.set_scheduler_timesteps(1000)
+
+            iteration = int(iteration / nsteps * 1000)
+            
+            positive_latents = diffuser.predict_noise(iteration, latents_steps[0], positive_text_embeddings, guidance_scale=1)
+            neutral_latents = diffuser.predict_noise(iteration, latents_steps[0], neutral_text_embeddings, guidance_scale=1)
+
+        with finetuner:
+            negative_latents = diffuser.predict_noise(iteration, latents_steps[0], positive_text_embeddings, guidance_scale=1)
+
+        positive_latents.requires_grad = False
+        neutral_latents.requires_grad = False
+
+        loss = criteria(negative_latents, neutral_latents - (negative_guidance*(positive_latents - neutral_latents))) #loss = criteria(e_n, e_0) works the best try 5000 epochs
+        
+        loss.backward()
+        optimizer.step()
+
+    torch.save(finetuner.state_dict(), save_path)
+
+    del diffuser, loss, optimizer, finetuner, negative_latents, neutral_latents, positive_latents, latents_steps, latents
+
+    torch.cuda.empty_cache()
+if __name__ == '__main__':
+
+    import argparse
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--prompt', required=True)
+    parser.add_argument('--modules', required=True)
+    parser.add_argument('--freeze_modules', nargs='+', required=True)
+    parser.add_argument('--save_path', required=True)
+    parser.add_argument('--iterations', type=int, required=True)
+    parser.add_argument('--lr', type=float, required=True)
+    parser.add_argument('--negative_guidance', type=float, required=True)
+
+    train(**vars(parser.parse_args()))

util.py

@@ -0,0 +1,107 @@
+from PIL import Image
+from matplotlib import pyplot as plt
+import textwrap
+
+
+def to_gif(images, path):
+
+    images[0].save(path, save_all=True,
+                   append_images=images[1:], loop=0, duration=len(images) * 20)
+
+
+def figure_to_image(figure):
+
+    figure.set_dpi(300)
+
+    figure.canvas.draw()
+
+    return Image.frombytes('RGB', figure.canvas.get_width_height(), figure.canvas.tostring_rgb())
+
+
+def image_grid(images, outpath=None, column_titles=None, row_titles=None):
+
+    n_rows = len(images)
+    n_cols = len(images[0])
+
+    fig, axs = plt.subplots(nrows=n_rows, ncols=n_cols,
+                            figsize=(n_cols, n_rows), squeeze=False)
+
+    for row, _images in enumerate(images):
+
+        for column, image in enumerate(_images):
+            ax = axs[row][column]
+            ax.imshow(image)
+            if column_titles and row == 0:
+                ax.set_title(textwrap.fill(
+                    column_titles[column], width=12), fontsize='x-small')
+            if row_titles and column == 0:
+                ax.set_ylabel(row_titles[row], rotation=0, fontsize='x-small', labelpad=1.6 * len(row_titles[row]))
+            ax.set_xticks([])
+            ax.set_yticks([])
+
+    plt.subplots_adjust(wspace=0, hspace=0)
+
+    if outpath is not None:
+        plt.savefig(outpath, bbox_inches='tight', dpi=300)
+        plt.close()
+    else:
+        plt.tight_layout(pad=0)
+        image = figure_to_image(plt.gcf())
+        plt.close()
+        return image
+
+
+
+
+    
+
+
+def get_module(module, module_name):
+
+    if isinstance(module_name, str):
+        module_name = module_name.split('.')
+
+    if len(module_name) == 0:
+        return module
+    else:
+        module = getattr(module, module_name[0])
+        return get_module(module, module_name[1:])
+
+
+def set_module(module, module_name, new_module):
+
+    if isinstance(module_name, str):
+        module_name = module_name.split('.')
+
+    if len(module_name) == 1:
+        return setattr(module, module_name[0], new_module)
+    else:
+        module = getattr(module, module_name[0])
+        return set_module(module, module_name[1:], new_module)
+
+
+def freeze(module):
+
+    for parameter in module.parameters():
+
+        parameter.requires_grad = False
+
+
+def unfreeze(module):
+
+    for parameter in module.parameters():
+
+        parameter.requires_grad = True
+
+
+def get_concat_h(im1, im2):
+    dst = Image.new('RGB', (im1.width + im2.width, im1.height))
+    dst.paste(im1, (0, 0))
+    dst.paste(im2, (im1.width, 0))
+    return dst
+
+def get_concat_v(im1, im2):
+    dst = Image.new('RGB', (im1.width, im1.height + im2.height))
+    dst.paste(im1, (0, 0))
+    dst.paste(im2, (0, im1.height))
+    return dst