concept ablation

.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,53 @@
+---
+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
+---
+
+
+
+# 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}
+}
+```

app.py

@@ -0,0 +1,233 @@
+import gradio as gr
+import torch
+import sys
+from pathlib import Path
+from trainer import train_submit, inference
+    
+import os
+model_map = {'Van Gogh' : 'models/vangogh_ablation_delta.bin', 
+             'Greg Rutkowski' : 'models/greg_rutkowski_ablation_delta.bin', 
+             }
+
+ORIGINAL_SPACE_ID = 'nupurkmr9/concept-ablation'
+SPACE_ID = os.getenv('SPACE_ID')
+
+SHARED_UI_WARNING = f'''## Attention - the demo requires at least 24GB VRAM for style and object removal, 24 GB VRAM for memorized image removal. Please 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>. This demo is partly adapted from https://huggingface.co/spaces/baulab/Erasing-Concepts-In-Diffusion.
+'''
+
+sys.path.append("concept-ablation-diffusers")
+
+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 [Concept Ablation](https://www.cs.cmu.edu/~concept-ablation/).  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 ablated and pre-trained models.  We have also provided several other pre-fine-tuned models with artistic styles and concepts ablated (Check out the "Ablated Model" drop-down). You can also train and run your own custom models. Check out the "train" section for custom ablation 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="Ablated Models",
+                                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="Ablated",
+                                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 ablate any concept from Stable Diffusion.   Enter the name of the concept and select the kind of concept (e.g. object, style, memorization). You will also need to select a parent anchor concept e.g. cats when ablating grumpy cat, painting when ablating an artists\' style. When ablating a specific object or memorized image, you also need to either provide OpenAI API key or upload a file with 50-200 prompts corresponding to the ablation concept. With default settings, it takes about 20 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/nupurkmr9/concept-ablation).')
+
+                with gr.Row():
+
+                    with gr.Column(scale=3):
+                        mem_impath = []
+
+                        self.prompt_input = gr.Text(
+                            placeholder="Enter concept to remove... e.g. van gogh",
+                            label="prompt",
+                            info="Name of the concept to ablate from Model"
+                        )
+                        self.anchor_prompt = gr.Text(
+                            placeholder="Enter anchor concept... e.g. painting",
+                            label="anchor prompt",
+                            info="Name of the anchor concept (superset of the concept to be ablated)"
+                        )
+
+                        choices = ['style', 'object', 'memorization']
+                        
+                        self.concept_type = gr.Dropdown(
+                            choices=choices,
+                            value='style',
+                            label='Ablated concept type',
+                            info='Ablated concept type'
+                        )
+
+                        self.reg_lambda = gr.Number(
+                            value=0,
+                            label="Regularization loss",
+                            info='Whether to add regularization loss on anchor concept. 1.0 when common words in ablated and anchor prompt e.g. grumpy cat and cat'
+                        )
+
+                        self.iterations_input = gr.Number(
+                            value=100,
+                            precision=0,
+                            label="Iterations",
+                            info='iterations used to train'
+                        )
+
+                        self.lr_input = gr.Number(
+                            value=2e-6,
+                            label="Learning Rate",
+                            info='Learning rate used to train'
+                        )
+
+                        visible_openai_key = True
+                        self.openai_key = gr.Text(
+                            placeholder="Enter openAI API key or atleast 50 prompts if concept type is object/memorization",
+                            label="OpenAI API key or Prompts (Required when concept type is object or memorization)",
+                            info="If concept type is object, we use chatGPT to generate a set of prompts correspondig to the ablation concept. If concept type is memorization, we use ChatGPT to generate paraphrases of the text prompt that generates memorized image. You can either provide the api key or a set of desired prompts (atleast 50). For reference please check example prompts at https://github.com/nupurkmr9/concept-ablation/blob/main/assets/finetune_prompts/ ",
+                            visible=visible_openai_key
+                        )
+
+                        visible = True
+                        mem_impath.append(gr.Files(label=f'''Upload the memorized image if concept type is memorization''', visible=visible))
+
+
+                    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.anchor_prompt, 
+            self.concept_type,
+            self.reg_lambda,
+            self.iterations_input,
+            self.lr_input,
+            self.openai_key,
+        ] + mem_impath,
+        outputs=[self.train_button,  self.train_status, self.download, self.model_dropdown]
+        )
+
+    def train(self, prompt, anchor_prompt, concept_type, reg_lambda, iterations, lr, openai_key, *inputs):
+        self.train_status.update(value='')
+        if self.training:
+            return [gr.update(interactive=True, value='Train'), gr.update(value='Someone else is training... Try again soon'), None, gr.update()]
+
+        randn = torch.randint(1, 10000000, (1,)).item()
+
+        save_path = f"models/{randn}_{prompt.lower().replace(' ', '')}"
+        os.makedirs(save_path, exist_ok=True)
+        self.training = True
+        mem_impath = inputs[:1]
+        train_submit(prompt, anchor_prompt, concept_type, reg_lambda, iterations, lr, openai_key, save_path, mem_impath)
+
+        self.training = False
+
+        torch.cuda.empty_cache()
+
+        modelpath = sorted(Path(save_path).glob('*.bin'))[0]
+        model_map[f"Custom_{prompt.lower().replace(' ', '')}"] = modelpath
+
+        return [gr.update(interactive=True, value='Train'), gr.update(value='Done Training! \n Try your custom model in the "Test" tab'), modelpath, 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
+        n_steps = 50
+
+        generator = torch.manual_seed(seed)
+
+        model_path = model_map[model_name]
+
+        torch.cuda.empty_cache()
+
+        generator = torch.manual_seed(seed)
+
+        orig_image, edited_image = inference(model_path, prompt, n_steps, generator)
+
+        torch.cuda.empty_cache()
+
+        return edited_image, orig_image
+
+
+demo = Demo()
+

assets/painting.txt

@@ -0,0 +1,200 @@
+Purple Paints
+Painting a Seascape in Water color part 1
+Why Arts & Crafts Are Beneficial to Early Childhood Education, Creve Coeur, Missouri
+Simple Cherry Blossom Tree Painting images
+Diy Wall Painting, Mural Painting, Paintings, Murals For Kids, Kids Wall Murals, Watercolor Walls, Mural Wall Art, Diy Canvas Art, Art Abstrait
+Watch Doris Allen paints Roses on porcelain, China Painter Doris Drew Allen
+Oil Painting with Kelsey May Connor in the tradition of Classical Realism | Art Classes Malta
+Peinture  LHuile  Fruits  Ltude  Youtube
+coffee painting canvas diy acrylic step by step for beginners #stepbysteppainting #tracie
+Sara Barcus paints Bob Ross
+Stock Video Footage of Little Baby Girl Paints Water Colors
+"Jill Baker working on her painting ""Lady In Waiting"" (2016)"
+Children painting a wall
+Close up shot of the process of mixing oil paint with palette knife Stock Footage
+toile : A young blonde woman draws a tree with paints Vidéos Libres De Droits
+Acrylic Paint Grass - How To Paint Grass - with Acrylic
+Snowfield Sentinel - En Plein Air Oil Demonstration
+Artist painting A Stock Photography
+peindre l 39 aquarelle youtube
+Monet inspired Water Lily Pond - good way to practice the art of shadows
+щетка для волос : A pretty woman with red hair, paints a picture on canvas, which stands on the easel. The lady is in the open air near the lake of the river, she draws from life Стоковые видеозаписи
+Wall Painting Techniques Youtube
+Painting under the pergola, Yeniköy 2014
+Practical Mom: Bleeding Crepe Paper Art Project for Little Kids
+painter man at work with paint brush, easel, canvas and...
+Watercolor landscape painting
+How to paint birds in watercolour
+Acrylic Painting Techniques - How to Paint Flowers - Parrot Tulip
+Painting cars toddler activity
+A painting student with her canvas at Dickinson Farm.
+Blonde woman in denim clothing finding creativity by painting a canvas hung on white wall.
+beautiful young woman painter at work,
+which : A pretty woman with red hair, paints a picture on canvas, which stands on the easel. The lady is in the open air near the lake of the river, she draws from life Stock Footage
+how-to-paint-a-crown-crested-crane-in-acrylic-background
+Slide Painting: A super fun indoor or outdoor process art activity for toddlers or preschoolers! Use cars, balls, or anything that rolls!
+Poppy Spree - Week 2 Painting
+How-To Make Colored Milk Explosion
+Abstract Acrylic Painting Demo - Watercolor Look - Color Explosion - Easy how to abstract
+Seine artist by triciamary
+Free Painting Lessons - Cherry Blossom Bridge - Commentary by Acrylic Artist Brandon Schaefer
+3 blue trees on ice - painting video step by step demonstration
+Martin Grealish painting in Venice.jpg
+Live Painter for Elegant Wedding
+Long-haired  artist paints on canvas
+Powe Painting Co Cover Photo
+Live Painting Part 1 (how to begin)
+which : A pretty woman with red hair, paints a picture on canvas, which stands on the easel. The lady is in the open air near the lake of the river, she draws from life Stock Footage
+Artist paints with oil paints
+Stock Video Footage of hand of a young woman painting in purple
+Watercolour tutorial - Summer shadows, cottages and flowers
+Scrape pastel with knife to release a snowfall!
+Acrylic pour painting process video
+Woman painting stock photo
+student painting at the Sainsbury Centre
+Honfleur Painter
+Bald Eagle in Wetlands Original Mini Painting on Easel
+painting in nature things to do in virginia beach
+Time lapse aquarelle painting Stock Footage
+Watercolour Aquarelle Poppies Poppy Painting Demo Youtube
+Wedding Artist Ben Keys of Wed on Canvas painting live at wedding reception.  //  The Graystone Inn, Wilmington, NC // Photo Courtesy of Blueberry Creative
+NIOS - 225 Painting - Guide Book For Class 10th - English Medium
+Fearless Watercolor #1 Frozen Marsh Full Watercolor Demonstration Music Ted Yoder
+Oil Painting with Kelsey May Connor in the tradition of Classical Realism | Art Classes Malta
+Michael Chambers, painting in process
+KNIFE PAINTING - WATER LILY by NELLY LESTRADE (english subtitles)
+A pretty woman with red hair, paints a picture on canvas, which stands on the easel. The lady is in the open air near the lake of the river, she draws from life
+Sip & Sketch at Acquiesce Winery!
+painting my bedroom wall youtube how to paint a galaxy wall mural in a spaceship themed
+Flat art painter workshop with paint supplies equipment tools ba Stock Illustration
+Gold Leaf Gilding in Westminster London
+Commission a portrait-artist-brisbane-perth-sydney-melbourne
+My very first brushstrokes as a mural painter
+Stock Video Footage of Little Baby Girl Paints Water Colors
+bing images of flower portraits watercolor by billy showell | VIDEO :: DVD - Watercolour Flower Portraits with Billy Showell
+Haidee-Jo Summers painting the winning painting at Paint Out Norwich 2014
+Artist paints a picture of oil paint brush in hand with palette Stock Footage
+img 9546 Sam Flores, recap: live painting in Golden Gate Park upper playground Sam Flores power to the peaceful paint live painting graffiti fifty24sf gallery art
+Artist at Michael Orwick (DragonFire Gallery) Stormy Weather Arts Festival Quick Draw event at Tolovana Inn
+Painting at the Tower - Stephen B Whatley, 2000 (Stephen B Whatley) Tags: uk windows england orange detail building london art architecture painting studio paint artist 2000 vibrant perspective millenium canvas creation painter expressionism expressionist toweroflondon palette damncool stkatharinedocks mywinners platinumphoto goldstaraward stephenbwhatley
+Video aula curso Pintura em tela com espatula part 1
+"getlinkyoutube.com-LIBERO! Full video ""orchidee"" dal artista Igor Sakharov"
+Weekly Specials Painting: Royal Paint By Number Mini Junior Dolphins
+Woman artist painting landscape outdoor Stock Footage
+painting styrofoam trays for a spring toddler art project
+Madinah Wilson, with the Biden Institute, signs a poster on the Green at the University of Delaware on Tuesday during an event for National Voter Registration Day that got students visiting the Green to register to vote.
+female hands painting a water landscape with oil paints using palette knife. 4k - cavalletto attrezzatura per arti e mestieri video stock e b–roll
+Art studio #miniature.  Dioramas and Clever Things    Dollhouse Miniatures and Accessories
+INDIANA - BOB ROSS EXHIBIT - 1-29-2021 (1).jpeg
+How to paint a CLOUDY sky tutorial - Jennings644
+Catherine Soucy
+Pin By Marionette Taboniar On Art Videos By Marionette
+coffee painting canvas diy acrylic step by step for beginners #stepbysteppainting #tracie
+Anita Jansen Youtube Watercolor Art Watercolor Trees Art
+Beautiful young woman painter at work, on grey background — Stock Photo
+In my natural habitat 🎨🎨 ~ I'm currently accepting tea recommendations. Tell me you're favourite teas(or hot drinks) in my stories! ☕️☕️ ~ #creatingsunflowers #mystudio #markmaking #naturallycurlyhair #curlybeauties #homeinthestudio #coloraddict #strongfit #fitnessenthusiast #collectart #originalpainting #colormehappy #workingprocess #intuitiveart #decorativeart #girlgains #worksonpaper #modernartist #womenwhopaint #girlgains
+Judith paints in a Bob Ross style.
+Royal & Langnickel: Painting by Number (Splish Splash)
+Close up of a needlework landscape being made at hong ngoc handicraft center Stock Footage
+Watercolour Tutorial Blea Tarn Cottage Lake District
+Squishy Painting from make-it-your-own.com (Crafts & activities for kids)
+Painter-girl en plein air
+Kids painting art outdoor activity, montessori homeschooling education
+HM Altered Book Project Page 3 - Art Journaling - Mixed Media - How to
+Fototapete - Talented male Artist Works on Abstract Oil Painting, with Broad Strokes of Paint Brush he Creates Modern Masterpiece. Dark and Messy Creative Studio where Large Canvas Stands on Easel Illuminated
+How I Start My Plein Air Paintings
+acrylic painting how to step by step flower painting mixed media flower heads 7 acrylic
+Oil painting classes  oil painting classes in Ameerpet  institution for oil painting classes in hyderabad  canvas painting classes in hyderabad  Oil paintings for sales  oil painting images  oil painting in canvas
+палитра : Young Beautiful Female Artist is in an Art Studio, Sitting Behind an Easel and Painting on Canvas. Drawing Process: in the Art Studio of the Artists Hand Art Girl with a Brush Painting on Canvas.4K Стоковые видеозаписи
+What is Plein Air Painting like in Ireland?
+Seven Tips for Setting up an Impromptu Garden Art Studio
+young woman painting - artist stock videos & royalty-free footage
+Pébéo - Mixed Media: mixing Studio Acrylics, Vitrail et Fantasy colors on a 3D frame - YouTube
+A painter To paint
+How to paint a pumpkin canvas, art skills not required!!!
+http://funkidos.com/pictures-world/art-world/hyper-realistic-paintings-on-wood
+Web Design for Oil Painters
+Painters palette with oil paints Stock Footage
+live-wedding-painting-by-ben-keys-live-event-artist-of-wed-on-canvas-garden-wedding-gibbes-museum-pure-luxe-bride-for-the-knot
+Lee's Painting logo
+Sir Winston Churchill painting
+The girl takes blue oil paint from the palette and spreads it on the canvas with Stock Footage
+Van Gogh Sunflowers | You Paint the Masters
+Painting lesson high school - tutorial de pintura para secundaria
+Crayon Art... now this is even cooler than the other kind of
+CoachingMasteryStudio-940-v1-png
+finished tree with pink flowers and within the circle on a gold background
+3 blue trees on ice - painting video step by step demonstration
+Talented male Artist Works on Abstract Oil Painting, with Broad Strokes of Paint Brush he Creates Modern Masterpiece. Dark and Messy Creative Studio where Large Canvas Stands on Easel Illuminated | Shutterstock HD Video #1036269983
+magic painting logo designs
+Workshop artistic, painting pigments color palette, color picker
+Toddler Art — Autumn Tree Painting
+Children painting pottery 2 Royalty Free Stock Photography
+4K Talented young graffiti street artist working on a mural in urban area Stock Footage
+fairytale castle mural for kids bedroom
+how to paint how to paint clouds in acrylic instructional painting
+A pretty woman with red hair, paints a picture on canvas, which stands on the easel. The lady is in the open air near the lake of the river, she draws from life
+The girl draws a picture by numbers,with yellow acrylic paints. Hobby for adults Stock Footage
+Sahara Mural (Jamila BC) Tags: cloud art sahara painting twilight mural desert dusk jamila jamilaproductions
+ressam : Close Up of Little Girl Drawing in Nature with Teacher. Art School Classes in the City Park. Slow Motion. Creativity Inspiration Expression Concept Stok Video
+Watercolor Tutorial Mountains | blue ridge mountains easy step by step watercolor
+Stock Video Footage of Street painting
+which : A pretty woman with red hair, paints a picture on canvas, which stands on the easel. The lady is in the open air near the lake of the river, she draws from life Stock Footage
+Palette,  canvas and easel Stock Image
+A woman paints a landscape with oil paints with a palette and brush. Slider shot Stock Footage
+1280x720 Blue Hydrangea Watercolor Doodle On Cotton Xuan Rice Paper By Amy
+Zero technique in watercolour but had loads of fun
+Taos - Plein Air Painter
+Artist drawing picture of a city landscape Stock Footage
+"""Painters and paintings"""
+Puff, Pass and Paint- 420-friendly painting in Orange County! 21+ tickets
+Melted Crayon Art | Unsimple Living
+Live Online Art Classes Webinar Oil Painting Sunset Pt1
+finger paint col.jpg
+Painting and wine class
+Artist Nancy Tankersley works on a plein air oil study of the landscape at King's Creek/Choptank Wetlands preserve.
+Using the palette knife to scrape away paint
+Close-up shot of a woman artist making brush strokes with oil on canvas. Impressionist painter working en plein air
+Students create paintings inspired by VanGogh's textured impasto
+Close-up shot of impressionist painter making final brush strokes during plein air painting
+Preschool Painting Activity Royalty Free Stock Images
+Advance Paint Pour Workshops
+Painting By Numbers Poppy Field
+Artist paints a picture of oil paint brush in hand with palette. Slider shot Stock Footage
+Art painting hobby leisure girl drawing picture. Art painting hobby. creative leisure. girl drawing a picture. talent inspiration creation and self expression stock photography
+1.5-Hour Children's Creative Art Session for 1 Child
+Painting the Nude in Watercolour Demonstration by Anthony Barrow BA MIFL
+The cherry blossoms in the Mt. Fuji Acrylic Painting Buying a new proven fact that everyone in your family will like doing together outside? Canvas Painting Tutorials, Easy Canvas Painting, Easy Paintings, Amazing Paintings, Mini Paintings, Pour Painting, Painting Videos, Painting Lessons, Small Canvas Art
+paintings in oil
+Printable Wall Murals japanese cherry blossoms by sole junkie youtube
+How to paint a Rose in watercolors, the impressionistic way
+Artist�s hand with paintbrush painting the picture          photo
+Childrens wall murals the muralist kids custom art for Equestrian wall mural
+St Patrick's Day Toddler Craft - Toilet paper shamrocks + pop up art
+Como pintar estrada, cerca e rvore - How to paint tree - how to paint landscape
+How to paint tree branches - Painting Tutorial
+spraying with a spray bottle
+Miniature paintings on the ground, a painting on the wall and a painting on an easel of a woman
+At the hairdresser on the market
+Paint With Kevin Hill Snow Covered Valley
+C. Rasmussen in her Los Angeles studio, 2018.
+Stock Video Footage of View from behind on child paints a sheet of paper by hand
+Sandra Gal seen through a ring flash tube as she works on her latest painting.
+Watercolour Art Classes
+Painting icon. Stock Footage
+Watercolour Trees - How To Paint A Tree In Watercolour
+Watercolor with Birgit O'Connor: Within the Flower
+London-Fine-Art-Studios-Painting.png
+Beautiful young woman painter at work, isolated on white
+Kindergarten and Day Care  Cartoon Wall Art Work In Mumbai
+Close-up of woman artist painting still life picture on canvas in art studio
+Modified Paint.Net logo
+Im Grand Canyon wird Kunst zelebriert. - Foto: Arizona Office of Tourism
+Melted crayons art
+Woman paints a picture outdoors.
+Snow Leopard Cub - Oils
+Stock Video Footage of close up of a brush that paints a canvas - painter
+Woman hand painting Stock Photo
+Watercolour Aquarelle Demo Poppies Coquelicots Aquarelle By
+painting 9 year 7 year pop painting american doll saige
+Contemporary man with smartphone taking photograph of the painting on easel

concept-ablation-diffusers/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 pix2pixzero
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

concept-ablation-diffusers/model_pipeline.py

@@ -0,0 +1,237 @@
+from typing import Callable, Optional
+
+import torch
+from accelerate.logging import get_logger
+from diffusers.models import AutoencoderKL, UNet2DConditionModel
+from diffusers.models.cross_attention import CrossAttention
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipeline
+from diffusers.pipelines.stable_diffusion.safety_checker import (
+    StableDiffusionSafetyChecker,
+)
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils.import_utils import is_xformers_available
+from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
+
+if is_xformers_available():
+    import xformers
+    import xformers.ops
+else:
+    xformers = None
+
+logger = get_logger(__name__)
+
+
+def set_use_memory_efficient_attention_xformers(
+    self, use_memory_efficient_attention_xformers: bool, attention_op: Optional[Callable] = None
+):
+    if use_memory_efficient_attention_xformers:
+        if self.added_kv_proj_dim is not None:
+            # TODO(Anton, Patrick, Suraj, William) - currently xformers doesn't work for UnCLIP
+            # which uses this type of cross attention ONLY because the attention mask of format
+            # [0, ..., -10.000, ..., 0, ...,] is not supported
+            raise NotImplementedError(
+                "Memory efficient attention with `xformers` is currently not supported when"
+                " `self.added_kv_proj_dim` is defined."
+            )
+        elif not is_xformers_available():
+            raise ModuleNotFoundError(
+                (
+                    "Refer to https://github.com/facebookresearch/xformers for more information on how to install"
+                    " xformers"
+                ),
+                name="xformers",
+            )
+        elif not torch.cuda.is_available():
+            raise ValueError(
+                "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is"
+                " only available for GPU "
+            )
+        else:
+            try:
+                # Make sure we can run the memory efficient attention
+                _ = xformers.ops.memory_efficient_attention(
+                    torch.randn((1, 2, 40), device="cuda"),
+                    torch.randn((1, 2, 40), device="cuda"),
+                    torch.randn((1, 2, 40), device="cuda"),
+                )
+            except Exception as e:
+                raise e
+
+        processor = CustomDiffusionXFormersAttnProcessor(
+            attention_op=attention_op)
+    else:
+        processor = CustomDiffusionAttnProcessor()
+
+    self.set_processor(processor)
+
+
+class CustomDiffusionAttnProcessor:
+    def __call__(
+        self,
+        attn: CrossAttention,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+    ):
+        batch_size, sequence_length, _ = hidden_states.shape
+        attention_mask = attn.prepare_attention_mask(
+            attention_mask, sequence_length, batch_size)
+        query = attn.to_q(hidden_states)
+
+        crossattn = False
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        else:
+            crossattn = True
+            if attn.cross_attention_norm:
+                encoder_hidden_states = attn.norm_cross(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+        if crossattn:
+            detach = torch.ones_like(key)
+            detach[:, :1, :] = detach[:, :1, :] * 0.
+            key = detach * key + (1 - detach) * key.detach()
+            value = detach * value + (1 - detach) * value.detach()
+
+        query = attn.head_to_batch_dim(query)
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        return hidden_states
+
+
+class CustomDiffusionXFormersAttnProcessor:
+    def __init__(self, attention_op: Optional[Callable] = None):
+        self.attention_op = attention_op
+
+    def __call__(self, attn: CrossAttention, hidden_states, encoder_hidden_states=None, attention_mask=None):
+        batch_size, sequence_length, _ = hidden_states.shape
+
+        attention_mask = attn.prepare_attention_mask(
+            attention_mask, sequence_length, batch_size)
+
+        query = attn.to_q(hidden_states)
+
+        crossattn = False
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        else:
+            crossattn = True
+            if attn.cross_attention_norm:
+                encoder_hidden_states = attn.norm_cross(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+        if crossattn:
+            detach = torch.ones_like(key)
+            detach[:, :1, :] = detach[:, :1, :] * 0.
+            key = detach * key + (1 - detach) * key.detach()
+            value = detach * value + (1 - detach) * value.detach()
+
+        query = attn.head_to_batch_dim(query).contiguous()
+        key = attn.head_to_batch_dim(key).contiguous()
+        value = attn.head_to_batch_dim(value).contiguous()
+
+        hidden_states = xformers.ops.memory_efficient_attention(
+            query, key, value, attn_bias=attention_mask, op=self.attention_op
+        )
+        hidden_states = hidden_states.to(query.dtype)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+        return hidden_states
+
+
+class CustomDiffusionPipeline(StableDiffusionPipeline):
+    r"""
+    Pipeline for custom diffusion model.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
+    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.).
+
+    Args:
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModel`]):
+            Frozen text-encoder. Stable Diffusion uses the text portion of
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
+            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
+        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
+        scheduler ([`SchedulerMixin`]):
+            A scheduler to be used in combination with `unet` to denoise the encoded image latents.
+        safety_checker ([`StableDiffusionSafetyChecker`]):
+            Classification module that estimates whether generated images could be considered offensive or harmful.
+            Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details.
+        feature_extractor ([`CLIPFeatureExtractor`]):
+            Model that extracts features from generated images to be used as inputs for the `safety_checker`.
+        modifier_token_id: list of id of tokens related to the target concept that are modified when ablated.
+    """
+    _optional_components = ["safety_checker",
+                            "feature_extractor", "modifier_token_id"]
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        unet: UNet2DConditionModel,
+        scheduler: SchedulerMixin,
+        safety_checker: StableDiffusionSafetyChecker,
+        feature_extractor: CLIPFeatureExtractor,
+        requires_safety_checker: bool = True,
+        modifier_token_id: list = [],
+    ):
+        super().__init__(vae,
+                         text_encoder,
+                         tokenizer,
+                         unet,
+                         scheduler,
+                         safety_checker,
+                         feature_extractor,
+                         requires_safety_checker)
+
+        self.modifier_token_id = modifier_token_id
+
+    def save_pretrained(self, save_path, parameter_group="cross-attn", all=False):
+        if all:
+            super().save_pretrained(save_path)
+        else:
+            delta_dict = {'unet': {}}
+            if parameter_group == 'embedding':
+                delta_dict['text_encoder'] = self.text_encoder.state_dict()
+            for name, params in self.unet.named_parameters():
+                if parameter_group == "cross-attn":
+                    if 'attn2.to_k' in name or 'attn2.to_v' in name:
+                        delta_dict['unet'][name] = params.cpu().clone()
+                elif parameter_group == "full-weight":
+                    delta_dict['unet'][name] = params.cpu().clone()
+                else:
+                    raise ValueError(
+                        "parameter_group argument only supports one of [cross-attn, full-weight, embedding]"
+                    )
+            torch.save(delta_dict, save_path)
+
+    def load_model(self, save_path):
+        st = torch.load(save_path)
+        print(st.keys())
+        if 'text_encoder' in st:
+            self.text_encoder.load_state_dict(st['text_encoder'])
+        for name, params in self.unet.named_parameters():
+            if name in st['unet']:
+                params.data.copy_(st['unet'][f'{name}'])

concept-ablation-diffusers/train.py

@@ -0,0 +1,1199 @@
+# This code is modified from the Huggingface repository: https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora.py, and
+import argparse
+import hashlib
+import itertools
+import json
+import logging
+import math
+import os
+import warnings
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+import transformers
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import ProjectConfiguration, set_seed
+from huggingface_hub import HfApi, create_repo
+from model_pipeline import (
+    CustomDiffusionAttnProcessor,
+    CustomDiffusionPipeline,
+    set_use_memory_efficient_attention_xformers,
+)
+from packaging import version
+from tqdm.auto import tqdm
+from transformers import AutoTokenizer, PretrainedConfig
+from utils import (
+    CustomDiffusionDataset,
+    PromptDataset,
+    collate_fn,
+    filter,
+    getanchorprompts,
+)
+
+import diffusers
+from diffusers import (
+    AutoencoderKL,
+    DDPMScheduler,
+    DiffusionPipeline,
+    DPMSolverMultistepScheduler,
+    UNet2DConditionModel,
+)
+from diffusers.models.cross_attention import CrossAttention
+from diffusers.optimization import get_scheduler
+from diffusers.utils import check_min_version, is_wandb_available
+from diffusers.utils.import_utils import is_xformers_available
+
+# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
+check_min_version("0.14.0")
+
+logger = get_logger(__name__)
+
+
+def create_custom_diffusion(unet, parameter_group):
+    for name, params in unet.named_parameters():
+        if parameter_group == "cross-attn":
+            if 'attn2.to_k' in name or 'attn2.to_v' in name:
+                params.requires_grad = True
+            else:
+                params.requires_grad = False
+        elif parameter_group == 'full-weight':
+            params.requires_grad = True
+        elif parameter_group == 'embedding':
+            params.requires_grad = False
+        else:
+            raise ValueError(
+                "parameter_group argument only cross-attn, full-weight, embedding"
+            )
+
+    # change attn class
+    def change_attn(unet):
+        for layer in unet.children():
+            if type(layer) == CrossAttention:
+                bound_method = set_use_memory_efficient_attention_xformers.__get__(
+                    layer, layer.__class__)
+                setattr(
+                    layer, 'set_use_memory_efficient_attention_xformers', bound_method)
+            else:
+                change_attn(layer)
+
+    change_attn(unet)
+    unet.set_attn_processor(CustomDiffusionAttnProcessor())
+    return unet
+
+
+def save_model_card(repo_id: str, images=None, base_model=str, prompt=str, repo_folder=None):
+    img_str = ""
+    for i, image in enumerate(images):
+        image.save(os.path.join(repo_folder, f"image_{i}.png"))
+        img_str += f"./image_{i}.png\n"
+
+    yaml = f"""
+        ---
+        license: creativeml-openrail-m
+        base_model: {base_model}
+        instance_prompt: {prompt}
+        tags:
+        - stable-diffusion
+        - stable-diffusion-diffusers
+        - text-to-image
+        - diffusers
+        - custom diffusion
+        inference: true
+        ---
+            """
+    model_card = f"""
+        # Custom Diffusion - {repo_id}
+
+        These are Custom Diffusion adaption weights for {base_model}. The weights were trained on {prompt} using [Custom Diffusion](https://www.cs.cmu.edu/~custom-diffusion). You can find some example images in the following. \n
+        {img_str[0]}
+        """
+    with open(os.path.join(repo_folder, "README.md"), "w") as f:
+        f.write(yaml + model_card)
+
+
+def import_model_class_from_model_name_or_path(pretrained_model_name_or_path: str, revision: str):
+    text_encoder_config = PretrainedConfig.from_pretrained(
+        pretrained_model_name_or_path,
+        subfolder="text_encoder",
+        revision=revision,
+    )
+    model_class = text_encoder_config.architectures[0]
+
+    if model_class == "CLIPTextModel":
+        from transformers import CLIPTextModel
+
+        return CLIPTextModel
+    elif model_class == "RobertaSeriesModelWithTransformation":
+        from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
+            RobertaSeriesModelWithTransformation,
+        )
+
+        return RobertaSeriesModelWithTransformation
+    else:
+        raise ValueError(f"{model_class} is not supported.")
+
+
+def freeze_params(params):
+    for param in params:
+        param.requires_grad = False
+
+
+def parse_args(input_args=None):
+    parser = argparse.ArgumentParser(
+        description="Simple example of a training script.")
+    parser.add_argument(
+        "--pretrained_model_name_or_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--revision",
+        type=str,
+        default=None,
+        required=False,
+        help="Revision of pretrained model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--tokenizer_name",
+        type=str,
+        default=None,
+        help="Pretrained tokenizer name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--concept_type",
+        type=str,
+        required=True,
+        choices=['style', 'object', 'memorization'],
+        help='the type of removed concepts'
+    )
+    parser.add_argument(
+        "--caption_target",
+        type=str,
+        required=True,
+        help="target style to remove, used when kldiv loss",
+    )
+    parser.add_argument(
+        "--instance_data_dir",
+        type=str,
+        default=None,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--class_data_dir",
+        type=str,
+        default=None,
+        help="A folder containing the training data of class images.",
+    )
+    parser.add_argument(
+        "--instance_prompt",
+        type=str,
+        help="The prompt with identifier specifying the instance",
+    )
+    parser.add_argument(
+        "--class_prompt",
+        type=str,
+        default=None,
+        help="The prompt to specify images in the same class as provided instance images.",
+    )
+    parser.add_argument(
+        "--mem_impath",
+        type=str,
+        default="",
+        help='the path to saved memorized image. Required when concept_type is memorization'
+    )
+    parser.add_argument(
+        "--validation_prompt",
+        type=str,
+        default=None,
+        help="A prompt that is used during validation to verify that the model is learning.",
+    )
+    parser.add_argument(
+        "--num_validation_images",
+        type=int,
+        default=2,
+        help="Number of images that should be generated during validation with `validation_prompt`.",
+    )
+    parser.add_argument(
+        "--validation_steps",
+        type=int,
+        default=500,
+        help=(
+            "Run dreambooth validation every X epochs. Dreambooth validation consists of running the prompt"
+            " `args.validation_prompt` multiple times: `args.num_validation_images`."
+        ),
+    )
+    parser.add_argument(
+        "--with_prior_preservation",
+        default=False,
+        action="store_true",
+        help="Flag to add prior preservation loss.",
+    )
+    parser.add_argument("--prior_loss_weight", type=float,
+                        default=1.0, help="The weight of prior preservation loss.")
+    parser.add_argument(
+        "--train_size",
+        type=int,
+        default=1000,
+        help='the number of generated images used for ablating the concept'
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="custom-diffusion-model",
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument(
+        "--num_class_images",
+        type=int,
+        default=1000,
+        help=(
+            "Minimal anchor class images. If there are not enough images already present in"
+            " class_data_dir, additional images will be sampled with class_prompt."
+        ),
+    )
+    parser.add_argument(
+        "--num_class_prompts",
+        type=int,
+        default=200,
+        help=(
+            "Minimal prompts used to generate anchor class images"
+        ),
+    )
+    parser.add_argument("--seed", type=int, default=42,
+                        help="A seed for reproducible training.")
+    parser.add_argument(
+        "--resolution",
+        type=int,
+        default=512,
+        help=(
+            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+            " resolution"
+        ),
+    )
+    parser.add_argument(
+        "--center_crop",
+        default=False,
+        action="store_true",
+        help=(
+            "Whether to center crop the input images to the resolution. If not set, the images will be randomly"
+            " cropped. The images will be resized to the resolution first before cropping."
+        ),
+    )
+    parser.add_argument(
+        "--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader."
+    )
+    parser.add_argument(
+        "--sample_batch_size", type=int, default=4, help="Batch size (per device) for sampling images."
+    )
+    parser.add_argument("--num_train_epochs", type=int, default=1)
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=int,
+        default=250,
+        help=(
+            "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
+            " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
+            " training using `--resume_from_checkpoint`."
+        ),
+    )
+    parser.add_argument(
+        "--checkpoints_total_limit",
+        type=int,
+        default=None,
+        help=(
+            "Max number of checkpoints to store. Passed as `total_limit` to the `Accelerator` `ProjectConfiguration`."
+            " See Accelerator::save_state https://huggingface.co/docs/accelerate/package_reference/accelerator#accelerate.Accelerator.save_state"
+            " for more docs"
+        ),
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help=(
+            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
+            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+        ),
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=1e-5,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument(
+        "--scale_lr",
+        action="store_true",
+        default=False,
+        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
+    )
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=2,
+        help=(
+            "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
+        ),
+    )
+    parser.add_argument(
+        "--parameter_group",
+        type=str,
+        default='cross-attn',
+        choices=['full-weight', 'cross-attn', 'embedding'],
+        help='parameter groups to finetune. Default: full-weight for memorization and cross-attn for others'
+    )
+    parser.add_argument(
+        "--loss_type_reverse",
+        type=str,
+        default='model-based',
+        help="loss type for reverse fine-tuning",
+    )
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="constant",
+        help=(
+            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+            ' "constant", "constant_with_warmup"]'
+        ),
+    )
+    parser.add_argument(
+        "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument(
+        "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
+    )
+    parser.add_argument("--adam_beta1", type=float, default=0.9,
+                        help="The beta1 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_beta2", type=float, default=0.999,
+                        help="The beta2 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_weight_decay", type=float,
+                        default=1e-2, help="Weight decay to use.")
+    parser.add_argument("--adam_epsilon", type=float, default=1e-08,
+                        help="Epsilon value for the Adam optimizer")
+    parser.add_argument("--max_grad_norm", default=1.0,
+                        type=float, help="Max gradient norm.")
+    parser.add_argument("--push_to_hub", action="store_true",
+                        help="Whether or not to push the model to the Hub.")
+    parser.add_argument("--hub_token", type=str, default=None,
+                        help="The token to use to push to the Model Hub.")
+    parser.add_argument(
+        "--hub_model_id",
+        type=str,
+        default=None,
+        help="The name of the repository to keep in sync with the local `output_dir`.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=(
+            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+        ),
+    )
+    parser.add_argument(
+        "--allow_tf32",
+        action="store_true",
+        help=(
+            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+        ),
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="tensorboard",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
+            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
+        ),
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default=None,
+        choices=["no", "fp16", "bf16"],
+        help=(
+            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+        ),
+    )
+    parser.add_argument(
+        "--prior_generation_precision",
+        type=str,
+        default=None,
+        choices=["no", "fp32", "fp16", "bf16"],
+        help=(
+            "Choose prior generation precision between fp32, fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+            " 1.10.and an Nvidia Ampere GPU.  Default to  fp16 if a GPU is available else fp32."
+        ),
+    )
+    parser.add_argument(
+        "--concepts_list",
+        type=str,
+        default=None,
+        help="Path to json containing multiple concepts, will overwrite parameters like instance_prompt, class_prompt, etc.",
+    )
+    parser.add_argument(
+        "--openai_key",
+        type=str,
+        default="",
+        help=(
+            "OPENAI API key. required for ablating objects and memorized images."
+        ),
+    )
+    parser.add_argument("--local_rank", type=int, default=-1,
+                        help="For distributed training: local_rank")
+    parser.add_argument(
+        "--enable_xformers_memory_efficient_attention", action="store_true", help="Whether or not to use xformers."
+    )
+    parser.add_argument("--hflip", action="store_true",
+                        help="Apply horizontal flip data augmentation.")
+    parser.add_argument("--noaug", action="store_true",
+                        help="Dont apply augmentation during data augmentation when this flag is enabled.")
+
+    if input_args is not None:
+        args = parser.parse_args(input_args)
+    else:
+        args = parser.parse_args()
+
+    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
+    if env_local_rank != -1 and env_local_rank != args.local_rank:
+        args.local_rank = env_local_rank
+
+    if args.with_prior_preservation:
+        if args.concepts_list is None:
+            if args.class_data_dir is None:
+                raise ValueError(
+                    "You must specify a data directory for class images.")
+            if args.class_prompt is None:
+                raise ValueError("You must specify prompt for class images.")
+    else:
+        # logger is not available yet
+        if args.class_data_dir is not None:
+            warnings.warn(
+                "You need not use --class_data_dir without --with_prior_preservation.")
+        if args.class_prompt is not None:
+            warnings.warn(
+                "You need not use --class_prompt without --with_prior_preservation.")
+
+    return args
+
+
+def main(args):
+    logging_dir = Path(args.output_dir, args.logging_dir)
+
+    accelerator_project_config = ProjectConfiguration(
+        total_limit=args.checkpoints_total_limit)
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        mixed_precision=args.mixed_precision,
+        log_with=args.report_to,
+        project_dir=logging_dir,
+        project_config=accelerator_project_config,
+    )
+
+    if args.report_to == "wandb":
+        if not is_wandb_available():
+            raise ImportError(
+                "Make sure to install wandb if you want to use it for logging during training.")
+        import wandb
+
+    # Currently, it's not possible to do gradient accumulation when training two models with accelerate.accumulate
+    # This will be enabled soon in accelerate. For now, we don't allow gradient accumulation when training two models.
+    # TODO (patil-suraj): Remove this check when gradient accumulation with two models is enabled in accelerate.
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state, main_process_only=False)
+    if accelerator.is_local_main_process:
+        transformers.utils.logging.set_verbosity_warning()
+        diffusers.utils.logging.set_verbosity_info()
+    else:
+        transformers.utils.logging.set_verbosity_error()
+        diffusers.utils.logging.set_verbosity_error()
+
+    # We need to initialize the trackers we use, and also store our configuration.
+    # The trackers initializes automatically on the main process.
+    if accelerator.is_main_process:
+        print(vars(args))
+        accelerator.init_trackers("custom-diffusion", config=vars(args))
+
+    # If passed along, set the training seed now.
+    if args.seed is not None:
+        set_seed(args.seed)
+    if args.concepts_list is None:
+        args.concepts_list = [
+            {
+                "instance_prompt": args.instance_prompt,
+                "class_prompt": args.class_prompt,
+                "instance_data_dir": args.instance_data_dir,
+                "class_data_dir": args.class_data_dir,
+                "caption_target": args.caption_target,
+            }
+        ]
+    else:
+        with open(args.concepts_list, "r") as f:
+            args.concepts_list = json.load(f)
+
+    # Generate class images if prior preservation is enabled.
+    for i, concept in enumerate(args.concepts_list):
+        # directly path to ablation images and its corresponding prompts is provided.
+        if (concept['instance_prompt'] is not None and concept['instance_data_dir'] is not None):
+            break
+
+        class_images_dir = Path(concept['class_data_dir'])
+        if not class_images_dir.exists():
+            class_images_dir.mkdir(parents=True, exist_ok=True)
+        os.makedirs(f'{class_images_dir}/images', exist_ok=True)
+
+        # we need to generate training images
+        if len(list(Path(os.path.join(class_images_dir, 'images')).iterdir())) < args.num_class_images:
+            torch_dtype = torch.float16 if accelerator.device.type == "cuda" else torch.float32
+            if args.prior_generation_precision == "fp32":
+                torch_dtype = torch.float32
+            elif args.prior_generation_precision == "fp16":
+                torch_dtype = torch.float16
+            elif args.prior_generation_precision == "bf16":
+                torch_dtype = torch.bfloat16
+            pipeline = DiffusionPipeline.from_pretrained(
+                args.pretrained_model_name_or_path,
+                torch_dtype=torch_dtype,
+                safety_checker=None,
+                revision=args.revision,
+            )
+            pipeline.scheduler = DPMSolverMultistepScheduler.from_config(
+                                    pipeline.scheduler.config)
+
+            pipeline.set_progress_bar_config(disable=True)
+            pipeline.to(accelerator.device)
+
+            # need to create prompts using class_prompt.
+            if not os.path.isfile(concept['class_prompt']):
+                # style based prompts are retrieved from laion dataset
+                if args.concept_type == 'style':
+                    with open(os.path.join(class_images_dir, 'painting.txt')) as f:
+                        class_prompt_collection = [
+                                x.strip() for x in f.readlines()]
+
+                # LLM based prompt collection.
+                else:
+                    class_prompt = concept['class_prompt']
+                    # in case of object query chatGPT to generate captions containing the anchor category
+                    if args.concept_type == 'object':
+                        class_prompt_collection, _ = getanchorprompts(
+                            pipeline, accelerator, class_prompt, args.concept_type, class_images_dir, args.openai_key, args.num_class_prompts)
+                        with open(class_images_dir / 'caption_anchor.txt', 'w') as f:
+                            for prompt in class_prompt_collection:
+                                f.write(prompt + '\n')
+                    # in case of memorization query chatGPT to generate different captions that can be paraphrase of the origianl caption
+                    elif args.concept_type == 'memorization':
+                        class_prompt_collection, caption_target = getanchorprompts(
+                            pipeline, accelerator, class_prompt, args.concept_type, class_images_dir, args.openai_key, args.num_class_prompts, mem_impath=args.mem_impath)
+                        concept['caption_target'] += f';*+{caption_target}'
+                        with open(class_images_dir / 'caption_target.txt', 'w') as f:
+                            f.write(concept['caption_target'])
+                        print(class_prompt_collection,
+                              concept['caption_target'])
+            # class_prompt is filepath to prompts.
+            else:
+                with open(concept['class_prompt']) as f:
+                    class_prompt_collection = [
+                        x.strip() for x in f.readlines()]
+
+            num_new_images = args.num_class_images
+            logger.info(
+                f"Number of class images to sample: {num_new_images}.")
+
+            sample_dataset = PromptDataset(
+                class_prompt_collection, num_new_images)
+            sample_dataloader = torch.utils.data.DataLoader(
+                sample_dataset, batch_size=args.sample_batch_size)
+
+            sample_dataloader = accelerator.prepare(sample_dataloader)
+
+            if os.path.exists(f'{class_images_dir}/caption.txt'):
+                os.remove(f'{class_images_dir}/caption.txt')
+            if os.path.exists(f'{class_images_dir}/images.txt'):
+                os.remove(f'{class_images_dir}/images.txt')
+
+            for example in tqdm(
+                sample_dataloader, desc="Generating class images", disable=not accelerator.is_local_main_process
+            ):
+                accelerator.wait_for_everyone()
+                with open(f'{class_images_dir}/caption.txt', 'a') as f1, open(f'{class_images_dir}/images.txt', 'a') as f2:
+                    images = pipeline(example["prompt"], num_inference_steps=25, guidance_scale=6., eta=1.).images
+
+                    for i, image in enumerate(images):
+                        hash_image = hashlib.sha1(
+                            image.tobytes()).hexdigest()
+                        image_filename = class_images_dir / \
+                            f"images/{example['index'][i]}-{hash_image}.jpg"
+                        image.save(image_filename)
+                        f2.write(str(image_filename)+'\n')
+                    f1.write('\n'.join(example["prompt"]) + '\n')
+                    accelerator.wait_for_everyone()
+
+            del pipeline
+
+        if args.concept_type == 'memorization':
+            filter(class_images_dir, args.mem_impath,
+                    outpath=str(class_images_dir / 'filtered'))
+            if os.path.exists(class_images_dir / 'caption_target.txt'):
+                with open(class_images_dir / 'caption_target.txt', 'r') as f:
+                    concept['caption_target'] = f.readlines()[0].strip()
+            class_images_dir = class_images_dir / 'filtered'
+
+        concept['class_prompt'] = os.path.join(
+            class_images_dir, 'caption.txt')
+        concept['class_data_dir'] = os.path.join(
+            class_images_dir, 'images.txt')
+        concept['instance_prompt'] = os.path.join(
+            class_images_dir, 'caption.txt')
+        concept['instance_data_dir'] = os.path.join(
+            class_images_dir, 'images.txt')
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    # Handle the repository creation
+    if accelerator.is_main_process:
+        if args.output_dir is not None:
+            os.makedirs(args.output_dir, exist_ok=True)
+
+        if args.push_to_hub:
+            print(args.hub_model_id or Path(args.output_dir).name)
+            repo_id = create_repo(
+                repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token
+            )
+            print(repo_id)
+            repo_id = args.hub_model_id
+
+    # Load the tokenizer
+    if args.tokenizer_name:
+        tokenizer = AutoTokenizer.from_pretrained(
+            args.tokenizer_name,
+            revision=args.revision,
+            use_fast=False,
+        )
+    elif args.pretrained_model_name_or_path:
+        tokenizer = AutoTokenizer.from_pretrained(
+            args.pretrained_model_name_or_path,
+            subfolder="tokenizer",
+            revision=args.revision,
+            use_fast=False,
+        )
+
+    # import correct text encoder class
+    text_encoder_cls = import_model_class_from_model_name_or_path(
+        args.pretrained_model_name_or_path, args.revision)
+
+    # Load scheduler and models
+    noise_scheduler = DDPMScheduler.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="scheduler")
+    text_encoder = text_encoder_cls.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
+    )
+    vae = AutoencoderKL.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision)
+    unet = UNet2DConditionModel.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
+    )
+
+    vae.requires_grad_(False)
+    if args.parameter_group != 'embedding':
+        text_encoder.requires_grad_(False)
+    unet = create_custom_diffusion(unet, args.parameter_group)
+
+    # For mixed precision training we cast the text_encoder and vae weights to half-precision
+    # as these models are only used for inference, keeping weights in full precision is not required.
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    # Move unet, vae and text_encoder to device and cast to weight_dtype
+    if accelerator.mixed_precision != "fp16":
+        unet.to(accelerator.device, dtype=weight_dtype)
+        text_encoder.to(accelerator.device, dtype=weight_dtype)
+    vae.to(accelerator.device, dtype=weight_dtype)
+
+    if args.enable_xformers_memory_efficient_attention:
+        if is_xformers_available():
+            import xformers
+            xformers_version = version.parse(xformers.__version__)
+            if xformers_version == version.parse("0.0.16"):
+                logger.warn(
+                    "xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details."
+                )
+            unet.enable_xformers_memory_efficient_attention()
+        else:
+            raise ValueError(
+                "xformers is not available. Make sure it is installed correctly")
+
+    if args.gradient_checkpointing:
+        unet.enable_gradient_checkpointing()
+        if args.parameter_group == 'embedding':
+            text_encoder.gradient_checkpointing_enable()
+    # Enable TF32 for faster training on Ampere GPUs,
+    # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
+    if args.allow_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+    if args.scale_lr:
+        args.learning_rate = (
+            args.learning_rate * args.gradient_accumulation_steps *
+            args.train_batch_size * accelerator.num_processes
+        )
+        if args.with_prior_preservation:
+            args.learning_rate = args.learning_rate * 2.
+
+    # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs
+    if args.use_8bit_adam:
+        try:
+            import bitsandbytes as bnb
+        except ImportError:
+            raise ImportError(
+                "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
+            )
+
+        optimizer_class = bnb.optim.AdamW8bit
+    else:
+        optimizer_class = torch.optim.AdamW
+
+    # Adding a modifier token which is optimized ####
+    # Code taken from https://github.com/huggingface/diffusers/blob/main/examples/textual_inversion/textual_inversion.py
+    modifier_token_id = []
+    if args.parameter_group == 'embedding':
+        assert args.concept_type != 'memorization', "embedding finetuning is not supported for memorization"
+
+        for concept in args.concept_list:
+            # Convert the caption_target to ids
+            token_ids = tokenizer.encode(
+                [concept['caption_target']], add_special_tokens=False)
+            print(token_ids)
+        # Check if initializer_token is a single token or a sequence of tokens
+        modifier_token_id += token_ids
+
+        # Freeze all parameters except for the token embeddings in text encoder
+        params_to_freeze = itertools.chain(
+            text_encoder.text_model.encoder.parameters(),
+            text_encoder.text_model.final_layer_norm.parameters(),
+            text_encoder.text_model.embeddings.position_embedding.parameters(),
+        )
+        freeze_params(params_to_freeze)
+        params_to_optimize = itertools.chain(
+            text_encoder.get_input_embeddings().parameters())
+    else:
+        if args.parameter_group == 'cross-attn':
+            params_to_optimize = itertools.chain([x[1] for x in unet.named_parameters() if (
+                'attn2.to_k' in x[0] or 'attn2.to_v' in x[0])])
+        if args.parameter_group == 'full-weight':
+            params_to_optimize = itertools.chain(unet.parameters())
+
+    # Optimizer creation
+    optimizer = optimizer_class(
+        params_to_optimize,
+        lr=args.learning_rate,
+        betas=(args.adam_beta1, args.adam_beta2),
+        weight_decay=args.adam_weight_decay,
+        eps=args.adam_epsilon,
+    )
+
+    # Dataset and DataLoaders creation:
+    train_dataset = CustomDiffusionDataset(
+        concepts_list=args.concepts_list,
+        concept_type=args.concept_type,
+        tokenizer=tokenizer,
+        with_prior_preservation=args.with_prior_preservation,
+        size=args.resolution,
+        center_crop=args.center_crop,
+        num_class_images=args.num_class_images,
+        hflip=args.hflip, aug=not args.noaug,
+    )
+
+    train_dataloader = torch.utils.data.DataLoader(
+        train_dataset,
+        batch_size=args.train_batch_size,
+        shuffle=True,
+        collate_fn=lambda examples: collate_fn(
+            examples, args.with_prior_preservation),
+        num_workers=args.dataloader_num_workers,
+    )
+
+    # Scheduler and math around the number of training steps.
+    overrode_max_train_steps = False
+    num_update_steps_per_epoch = math.ceil(
+        len(train_dataloader) / args.gradient_accumulation_steps)
+    if args.max_train_steps is None:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+        overrode_max_train_steps = True
+
+    lr_scheduler = get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
+        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
+    )
+
+    # Prepare everything with our `accelerator`.
+    if args.parameter_group == 'embedding':
+        text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+            text_encoder, optimizer, train_dataloader, lr_scheduler
+        )
+    else:
+        unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+            unet, optimizer, train_dataloader, lr_scheduler
+        )
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(
+        len(train_dataloader) / args.gradient_accumulation_steps)
+    if overrode_max_train_steps:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+    # Afterwards we recalculate our number of training epochs
+    args.num_train_epochs = math.ceil(
+        args.max_train_steps / num_update_steps_per_epoch)
+
+    # Train!
+    total_batch_size = args.train_batch_size * \
+        accelerator.num_processes * args.gradient_accumulation_steps
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {len(train_dataset)}")
+    logger.info(f"  Num batches each epoch = {len(train_dataloader)}")
+    logger.info(f"  Num Epochs = {args.num_train_epochs}")
+    logger.info(
+        f"  Instantaneous batch size per device = {args.train_batch_size}")
+    logger.info(
+        f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(
+        f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+    global_step = 0
+    first_epoch = 0
+
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint != "latest":
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the mos recent checkpoint
+            dirs = os.listdir(args.output_dir)
+            dirs = [d for d in dirs if d.startswith("checkpoint")]
+            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+            path = dirs[-1] if len(dirs) > 0 else None
+
+        if path is None:
+            accelerator.print(
+                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
+            )
+            args.resume_from_checkpoint = None
+        else:
+            accelerator.print(f"Resuming from checkpoint {path}")
+            accelerator.load_state(os.path.join(args.output_dir, path))
+            global_step = int(path.split("-")[1])
+
+            resume_global_step = global_step * args.gradient_accumulation_steps
+            first_epoch = global_step // num_update_steps_per_epoch
+            resume_step = resume_global_step % (
+                num_update_steps_per_epoch * args.gradient_accumulation_steps)
+
+    # Only show the progress bar once on each machine.
+    progress_bar = tqdm(range(global_step, args.max_train_steps),
+                        disable=not accelerator.is_local_main_process)
+    progress_bar.set_description("Steps")
+
+    for epoch in range(first_epoch, args.num_train_epochs):
+        if args.parameter_group == 'embedding':
+            text_encoder.train()
+        else:
+            unet.train()
+        for step, batch in enumerate(train_dataloader):
+            # Skip steps until we reach the resumed step
+            if args.resume_from_checkpoint and epoch == first_epoch and step < resume_step:
+                if step % args.gradient_accumulation_steps == 0:
+                    progress_bar.update(1)
+                continue
+
+            with accelerator.accumulate(unet) if args.parameter_group != 'embedding' else accelerator.accumulate(text_encoder):
+                # Convert images to latent space
+                latents = vae.encode(batch["pixel_values"].to(
+                    dtype=weight_dtype)).latent_dist.sample()
+                latents = latents * vae.config.scaling_factor
+
+                # Sample noise that we'll add to the latents
+                noise = torch.randn_like(latents)
+                bsz = latents.shape[0]
+                # Sample a random timestep for each image
+                timesteps = torch.randint(
+                    0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
+                timesteps = timesteps.long()
+
+                # Add noise to the latents according to the noise magnitude at each timestep
+                # (this is the forward diffusion process)
+                noisy_latents = noise_scheduler.add_noise(
+                    latents, noise, timesteps)
+
+                # Get the text embedding for conditioning
+                encoder_hidden_states = text_encoder(batch["input_ids"])[0]
+                encoder_anchor_hidden_states = text_encoder(
+                    batch["input_anchor_ids"])[0]
+
+                # Predict the noise residual
+                model_pred = unet(noisy_latents, timesteps,
+                                  encoder_hidden_states).sample
+                with torch.no_grad():
+                    model_pred_anchor = unet(noisy_latents[:encoder_anchor_hidden_states.size(
+                        0)], timesteps[:encoder_anchor_hidden_states.size(0)], encoder_anchor_hidden_states).sample
+
+                # Get the target for loss depending on the prediction type
+                if args.loss_type_reverse == 'model-based':
+                    if args.with_prior_preservation:
+                        target_prior = torch.chunk(noise, 2, dim=0)[1]
+                    target = model_pred_anchor
+                else:
+                    if noise_scheduler.config.prediction_type == "epsilon":
+                        target = noise
+                    elif noise_scheduler.config.prediction_type == "v_prediction":
+                        target = noise_scheduler.get_velocity(
+                            latents, noise, timesteps)
+                    else:
+                        raise ValueError(
+                            f"Unknown prediction type {noise_scheduler.config.prediction_type}")
+                    if args.with_prior_preservation:
+                        target, target_prior = torch.chunk(target, 2, dim=0)
+
+                if args.with_prior_preservation:
+                    # Chunk the noise and model_pred into two parts and compute the loss on each part separately.
+                    model_pred, model_pred_prior = torch.chunk(
+                        model_pred, 2, dim=0)
+                    mask = torch.chunk(batch["mask"], 2, dim=0)[0]
+                    # Compute instance loss
+                    loss = F.mse_loss(model_pred.float(),
+                                      target.float(), reduction="none")
+                    loss = (
+                        (loss * mask).sum([1, 2, 3]) / mask.sum([1, 2, 3])).mean()
+
+                    # Compute prior loss
+                    prior_loss = F.mse_loss(
+                        model_pred_prior.float(), target_prior.float(), reduction="mean")
+
+                    # Add the prior loss to the instance loss.
+                    loss = loss + args.prior_loss_weight * prior_loss
+                else:
+                    mask = batch["mask"]
+                    loss = F.mse_loss(model_pred.float(),
+                                      target.float(), reduction="none")
+                    loss = (
+                        (loss * mask).sum([1, 2, 3]) / mask.sum([1, 2, 3])).mean()
+
+                accelerator.backward(loss)
+                # Zero out the gradients for all token embeddings except the newly added
+                # embeddings for the concept, as we only want to optimize the concept embeddings
+                if args.parameter_group == 'embedding':
+                    if accelerator.num_processes > 1:
+                        grads_text_encoder = text_encoder.module.get_input_embeddings().weight.grad
+                    else:
+                        grads_text_encoder = text_encoder.get_input_embeddings().weight.grad
+                    # Get the index for tokens that we want to zero the grads for
+                    index_grads_to_zero = torch.arange(
+                        len(tokenizer)) != modifier_token_id[0]
+                    for i in range(len(modifier_token_id[1:])):
+                        index_grads_to_zero = index_grads_to_zero & (
+                            torch.arange(len(tokenizer)) != modifier_token_id[i])
+                    grads_text_encoder.data[index_grads_to_zero,
+                                            :] = grads_text_encoder.data[index_grads_to_zero, :].fill_(0)
+
+                if accelerator.sync_gradients:
+                    params_to_clip = (
+                        itertools.chain(text_encoder.parameters())
+                        if args.parameter_group == 'embedding'
+                        else itertools.chain([x[1] for x in unet.named_parameters() if ('attn2' in x[0])])
+                        if args.parameter_group == 'cross-attn'
+                        else itertools.chain(unet.parameters())
+                    )
+                    accelerator.clip_grad_norm_(
+                        params_to_clip, args.max_grad_norm)
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+            # Checks if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                progress_bar.update(1)
+                global_step += 1
+
+                if global_step % args.checkpointing_steps == 0:
+                    if accelerator.is_main_process:
+                        pipeline = CustomDiffusionPipeline.from_pretrained(
+                            args.pretrained_model_name_or_path,
+                            unet=accelerator.unwrap_model(unet),
+                            text_encoder=accelerator.unwrap_model(
+                                text_encoder),
+                            tokenizer=tokenizer,
+                            revision=args.revision,
+                            modifier_token_id=modifier_token_id,
+                        )
+                        save_path = os.path.join(
+                            args.output_dir, f"delta-{global_step}")
+                        pipeline.save_pretrained(
+                            save_path, parameter_group=args.parameter_group)
+                        logger.info(f"Saved state to {save_path}")
+
+            logs = {"loss": loss.detach().item(
+            ), "lr": lr_scheduler.get_last_lr()[0]}
+            progress_bar.set_postfix(**logs)
+            accelerator.log(logs, step=global_step)
+
+            if global_step >= args.max_train_steps:
+                break
+
+        if accelerator.is_main_process:
+            if args.validation_prompt is not None and global_step % args.validation_steps == 0:
+                logger.info(
+                    f"Running validation... \n Generating {args.num_validation_images} images with prompt:"
+                    f" {args.validation_prompt}."
+                )
+                # create pipeline
+                pipeline = CustomDiffusionPipeline.from_pretrained(
+                    args.pretrained_model_name_or_path,
+                    unet=accelerator.unwrap_model(unet),
+                    text_encoder=accelerator.unwrap_model(text_encoder),
+                    tokenizer=tokenizer,
+                    revision=args.revision,
+                    modifier_token_id=modifier_token_id,
+                )
+                pipeline.scheduler = DPMSolverMultistepScheduler.from_config(
+                    pipeline.scheduler.config)
+                pipeline = pipeline.to(accelerator.device)
+                pipeline.set_progress_bar_config(disable=True)
+
+                # run inference
+                generator = torch.Generator(
+                    device=accelerator.device).manual_seed(args.seed)
+                images = [
+                    pipeline(args.validation_prompt, num_inference_steps=25,
+                             generator=generator, eta=1.).images[0]
+                    for _ in range(args.num_validation_images)
+                ]
+
+                for tracker in accelerator.trackers:
+                    if tracker.name == "tensorboard":
+                        np_images = np.stack([np.asarray(img)
+                                             for img in images])
+                        tracker.writer.add_images(
+                            "validation", np_images, epoch, dataformats="NHWC")
+                    if tracker.name == "wandb":
+                        tracker.log(
+                            {
+                                "validation": [
+                                    wandb.Image(
+                                        image, caption=f"{i}: {args.validation_prompt}")
+                                    for i, image in enumerate(images)
+                                ]
+                            }
+                        )
+
+                del pipeline
+                torch.cuda.empty_cache()
+
+    accelerator.wait_for_everyone()
+    if accelerator.is_main_process:
+        unet = unet.to(torch.float32)
+        pipeline = CustomDiffusionPipeline.from_pretrained(
+            args.pretrained_model_name_or_path,
+            unet=accelerator.unwrap_model(unet),
+            text_encoder=accelerator.unwrap_model(text_encoder),
+            tokenizer=tokenizer,
+            revision=args.revision,
+            modifier_token_id=modifier_token_id,
+        )
+        save_path = os.path.join(args.output_dir, "delta.bin")
+        pipeline.save_pretrained(
+            save_path, parameter_group=args.parameter_group)
+
+        # run inference
+        if args.validation_prompt and args.num_validation_images > 0:
+            pipeline.scheduler = DPMSolverMultistepScheduler.from_config(
+                pipeline.scheduler.config)
+            pipeline = pipeline.to(accelerator.device)
+            pipeline.set_progress_bar_config(disable=True)
+
+            # run inference
+            generator = torch.Generator(
+                device=accelerator.device).manual_seed(args.seed)
+            images = [
+                pipeline(args.validation_prompt, num_inference_steps=25,
+                         generator=generator, eta=1.).images[0]
+                for _ in range(args.num_validation_images)
+            ]
+
+            for tracker in accelerator.trackers:
+                if tracker.name == "tensorboard":
+                    np_images = np.stack([np.asarray(img) for img in images])
+                    tracker.writer.add_images(
+                        "test", np_images, epoch, dataformats="NHWC")
+                if tracker.name == "wandb":
+                    tracker.log(
+                        {
+                            "test": [
+                                wandb.Image(
+                                    image, caption=f"{i}: {args.validation_prompt}")
+                                for i, image in enumerate(images)
+                            ]
+                        }
+                    )
+
+        if args.push_to_hub:
+            save_model_card(
+                repo_id,
+                images=images,
+                base_model=args.pretrained_model_name_or_path,
+                prompt=args.instance_prompt,
+                repo_folder=args.output_dir,
+            )
+            api = HfApi(token=args.hub_token)
+            api.upload_folder(
+                repo_id=repo_id,
+                folder_path=args.output_dir,
+                path_in_repo='.',
+                repo_type='model'
+            )
+
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

concept-ablation-diffusers/utils.py

@@ -0,0 +1,443 @@
+import os
+import random
+import shutil
+from io import BytesIO
+from pathlib import Path
+
+import numpy as np
+import openai
+import regex as re
+import requests
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+from torchvision import transforms
+from tqdm.auto import tqdm
+
+from diffusers import DPMSolverMultistepScheduler
+
+normalize = transforms.Normalize(
+    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225],
+)
+small_288 = transforms.Compose([
+    transforms.Resize(288),
+    transforms.ToTensor(),
+    normalize,
+])
+
+def collate_fn(examples, with_prior_preservation):
+    input_ids = [example["instance_prompt_ids"] for example in examples]
+    input_anchor_ids = [example["instance_anchor_prompt_ids"]
+                        for example in examples]
+    pixel_values = [example["instance_images"] for example in examples]
+    mask = [example["mask"] for example in examples]
+    # Concat class and instance examples for prior preservation.
+    # We do this to avoid doing two forward passes.
+    if with_prior_preservation:
+        input_ids += [example["class_prompt_ids"] for example in examples]
+        pixel_values += [example["class_images"] for example in examples]
+        mask += [example["class_mask"] for example in examples]
+
+    input_ids = torch.cat(input_ids, dim=0)
+    input_anchor_ids = torch.cat(input_anchor_ids, dim=0)
+    pixel_values = torch.stack(pixel_values)
+    mask = torch.stack(mask)
+    pixel_values = pixel_values.to(
+        memory_format=torch.contiguous_format).float()
+    mask = mask.to(memory_format=torch.contiguous_format).float()
+
+    batch = {
+        "input_ids": input_ids,
+        "input_anchor_ids": input_anchor_ids,
+        "pixel_values": pixel_values,
+        "mask": mask.unsqueeze(1)
+    }
+    return batch
+
+
+class PromptDataset(Dataset):
+    "A simple dataset to prepare the prompts to generate class images on multiple GPUs."
+
+    def __init__(self, prompt, num_samples):
+        self.prompt = prompt
+        self.num_samples = num_samples
+
+    def __len__(self):
+        return self.num_samples
+
+    def __getitem__(self, index):
+        example = {}
+        example["prompt"] = self.prompt[index % len(self.prompt)]
+        example["index"] = index
+        return example
+
+
+class CustomDiffusionDataset(Dataset):
+    """
+    A dataset to prepare the instance and class images with the prompts for fine-tuning the model.
+    It pre-processes the images and the tokenizes prompts.
+    """
+
+    def __init__(
+        self,
+        concepts_list,
+        concept_type,
+        tokenizer,
+        size=512,
+        center_crop=False,
+        with_prior_preservation=False,
+        num_class_images=200,
+        hflip=False,
+        aug=True,
+    ):
+        self.size = size
+        self.center_crop = center_crop
+        self.tokenizer = tokenizer
+        self.interpolation = Image.BILINEAR
+        self.aug = aug
+        self.concept_type = concept_type
+
+        self.instance_images_path = []
+        self.class_images_path = []
+        self.with_prior_preservation = with_prior_preservation
+        for concept in concepts_list:
+            with open(concept["instance_data_dir"], "r") as f:
+                inst_images_path = f.read().splitlines()
+            with open(concept["instance_prompt"], "r") as f:
+                inst_prompt = f.read().splitlines()
+            inst_img_path = [(x, y, concept['caption_target'])
+                             for (x, y) in zip(inst_images_path, inst_prompt)]
+            self.instance_images_path.extend(inst_img_path)
+
+            if with_prior_preservation:
+                class_data_root = Path(concept["class_data_dir"])
+                if os.path.isdir(class_data_root):
+                    class_images_path = list(class_data_root.iterdir())
+                    class_prompt = [concept["class_prompt"]
+                                    for _ in range(len(class_images_path))]
+                else:
+                    with open(class_data_root, "r") as f:
+                        class_images_path = f.read().splitlines()
+                    with open(concept["class_prompt"], "r") as f:
+                        class_prompt = f.read().splitlines()
+
+                class_img_path = [(x, y) for (x, y) in zip(
+                    class_images_path, class_prompt)]
+                self.class_images_path.extend(
+                    class_img_path[:num_class_images])
+
+        random.shuffle(self.instance_images_path)
+        self.num_instance_images = len(self.instance_images_path)
+        self.num_class_images = len(self.class_images_path)
+        self._length = max(self.num_class_images, self.num_instance_images)
+        self.flip = transforms.RandomHorizontalFlip(0.5 * hflip)
+
+        self.image_transforms = transforms.Compose(
+            [
+                self.flip,
+                transforms.Resize(
+                    size, interpolation=transforms.InterpolationMode.BILINEAR),
+                transforms.CenterCrop(
+                    size) if center_crop else transforms.RandomCrop(size),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5], [0.5]),
+            ]
+        )
+
+    def __len__(self):
+        return self._length
+
+    def preprocess(self, image, scale, resample):
+        outer, inner = self.size, scale
+        if scale > self.size:
+            outer, inner = scale, self.size
+        top, left = np.random.randint(
+            0, outer - inner + 1), np.random.randint(0, outer - inner + 1)
+        image = image.resize((scale, scale), resample=resample)
+        image = np.array(image).astype(np.uint8)
+        image = (image / 127.5 - 1.0).astype(np.float32)
+        instance_image = np.zeros((self.size, self.size, 3), dtype=np.float32)
+        mask = np.zeros((self.size // 8, self.size // 8))
+        if scale > self.size:
+            instance_image = image[top: top + inner, left: left + inner, :]
+            mask = np.ones((self.size // 8, self.size // 8))
+        else:
+            instance_image[top: top + inner, left: left + inner, :] = image
+            mask[top // 8 + 1: (top + scale) // 8 - 1, left //
+                 8 + 1: (left + scale) // 8 - 1] = 1.
+        return instance_image, mask
+
+    def __getprompt__(self, instance_prompt, instance_target):
+        if self.concept_type == 'style':
+            r = np.random.choice([0, 1, 2])
+            instance_prompt = f'{instance_prompt}, in the style of {instance_target}' if r == 0 else f'in {instance_target}\'s style, {instance_prompt}' if r == 1 else f'in {instance_target}\'s style, {instance_prompt}'
+        elif self.concept_type == 'object':
+            anchor, target = instance_target.split('+')
+            instance_prompt = instance_prompt.replace(anchor, target)
+        elif self.concept_type == 'memorization':
+            instance_prompt = instance_target.split('+')[1]
+        return instance_prompt
+
+    def __getitem__(self, index):
+        example = {}
+        instance_image, instance_prompt, instance_target = self.instance_images_path[
+            index % self.num_instance_images]
+        instance_image = Image.open(instance_image)
+        if not instance_image.mode == "RGB":
+            instance_image = instance_image.convert("RGB")
+        instance_image = self.flip(instance_image)
+        # modify instance prompt according to the concept_type to include target concept
+        # multiple style/object fine-tuning
+        if ';' in instance_target:
+            instance_target = instance_target.split(';')
+            instance_target = instance_target[index % len(instance_target)]
+
+        instance_anchor_prompt = instance_prompt
+        instance_prompt = self.__getprompt__(instance_prompt, instance_target)
+        # apply resize augmentation and create a valid image region mask
+        random_scale = self.size
+        if self.aug:
+            random_scale = np.random.randint(self.size // 3, self.size + 1) if np.random.uniform(
+            ) < 0.66 else np.random.randint(int(1.2 * self.size), int(1.4 * self.size))
+        instance_image, mask = self.preprocess(
+            instance_image, random_scale, self.interpolation)
+
+        if random_scale < 0.6 * self.size:
+            instance_prompt = np.random.choice(
+                ["a far away ", "very small "]) + instance_prompt
+        elif random_scale > self.size:
+            instance_prompt = np.random.choice(
+                ["zoomed in ", "close up "]) + instance_prompt
+
+        example["instance_images"] = torch.from_numpy(
+            instance_image).permute(2, 0, 1)
+        example["mask"] = torch.from_numpy(mask)
+
+        example["instance_prompt_ids"] = self.tokenizer(
+            instance_prompt,
+            truncation=True,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            return_tensors="pt",
+        ).input_ids
+        example["instance_anchor_prompt_ids"] = self.tokenizer(
+            instance_anchor_prompt,
+            truncation=True,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            return_tensors="pt",
+        ).input_ids
+
+        if self.with_prior_preservation:
+            class_image, class_prompt = self.class_images_path[index %
+                                                               self.num_class_images]
+            class_image = Image.open(class_image)
+            if not class_image.mode == "RGB":
+                class_image = class_image.convert("RGB")
+            example["class_images"] = self.image_transforms(class_image)
+            example["class_mask"] = torch.ones_like(example["mask"])
+            example["class_prompt_ids"] = self.tokenizer(
+                class_prompt,
+                truncation=True,
+                padding="max_length",
+                max_length=self.tokenizer.model_max_length,
+                return_tensors="pt",
+            ).input_ids
+
+        return example
+
+
+def isimage(path):
+    if 'png' in path.lower() or 'jpg' in path.lower() or 'jpeg' in path.lower():
+        return True
+
+
+def filter(folder, impath, outpath=None, unfiltered_path=None, threshold=0.15,
+           image_threshold=0.5, anchor_size=10, target_size=3, return_score=False):
+    model = torch.jit.load(
+        "./assets/sscd_imagenet_mixup.torchscript.pt")
+    if isinstance(folder, list):
+        image_paths = folder
+        image_captions = ["None" for _ in range(len(image_paths))]
+    elif Path(folder / 'images.txt').exists():
+        with open(f'{folder}/images.txt', "r") as f:
+            image_paths = f.read().splitlines()
+        with open(f'{folder}/caption.txt', "r") as f:
+            image_captions = f.read().splitlines()
+    else:
+        image_paths = [os.path.join(str(folder), file_path)
+                       for file_path in os.listdir(folder) if isimage(file_path)]
+        image_captions = ["None" for _ in range(len(image_paths))]
+
+    batch = small_288(Image.open(impath).convert('RGB')).unsqueeze(0)
+    embedding_target = model(batch)[0, :]
+
+    filtered_paths = []
+    filtered_captions = []
+    unfiltered_paths = []
+    unfiltered_captions = []
+    count_dict = {}
+    for im, c in zip(image_paths, image_captions):
+        if c not in count_dict:
+            count_dict[c] = 0
+        if isinstance(folder, list):
+            batch = small_288(im).unsqueeze(0)
+        else:
+            batch = small_288(Image.open(im).convert('RGB')).unsqueeze(0)
+        embedding = model(batch)[0, :]
+
+        diff_sscd = (embedding * embedding_target).sum()
+
+        if diff_sscd <= image_threshold:
+            filtered_paths.append(im)
+            filtered_captions.append(c)
+            count_dict[c] += 1
+        else:
+            unfiltered_paths.append(im)
+            unfiltered_captions.append(c)
+
+    # only return score
+    if return_score:
+        score = len(unfiltered_paths) / \
+            (len(unfiltered_paths)+len(filtered_paths))
+        return score
+
+    os.makedirs(outpath, exist_ok=True)
+    os.makedirs(f'{outpath}/samples', exist_ok=True)
+    with open(f'{outpath}/caption.txt', 'w') as f:
+        for each in filtered_captions:
+            f.write(each.strip() + '\n')
+
+    with open(f'{outpath}/images.txt', 'w') as f:
+        for each in filtered_paths:
+            f.write(each.strip() + '\n')
+            imbase = Path(each).name
+            shutil.copy(each, f'{outpath}/samples/{imbase}')
+
+    print('++++++++++++++++++++++++++++++++++++++++++++++++')
+    print('+ Filter Summary +')
+    print(f'+ Remained images: {len(filtered_paths)}')
+    print(f'+ Filtered images: {len(unfiltered_paths)}')
+    print('++++++++++++++++++++++++++++++++++++++++++++++++')
+
+    sorted_list = sorted(list(count_dict.items()),
+                         key=lambda x: x[1], reverse=True)
+    anchor_prompts = [c[0] for c in sorted_list[:anchor_size]]
+    target_prompts = [c[0] for c in sorted_list[-target_size:]]
+    return anchor_prompts, target_prompts, len(filtered_paths)
+
+
+def getanchorprompts(pipeline, accelerator, class_prompt, concept_type, class_images_dir, api_key, num_class_images=200, mem_impath=None):
+    openai.api_key = api_key
+    class_prompt_collection = []
+    caption_target = []
+    if concept_type == 'object':
+        messages = [{"role": "system", "content": "You can describe any image via text and provide captions for wide variety of images that is possible to generate."}]
+        messages = [{"role": "user", "content": f"Generate {num_class_images} captions for images containing a {class_prompt}. The caption should also contain the word \"{class_prompt}\" "}]
+        while True:
+            completion = openai.ChatCompletion.create(
+                model="gpt-3.5-turbo",
+                messages=messages
+            )
+            class_prompt_collection += [x for x in completion.choices[0].message.content.lower(
+            ).split('\n') if class_prompt in x]
+            messages.append(
+                {"role": "assistant", "content": completion.choices[0].message.content})
+            messages.append(
+                {"role": "user", "content": f"Generate {num_class_images-len(class_prompt_collection)} more captions"})
+            if len(class_prompt_collection) >= num_class_images:
+                break
+        class_prompt_collection = clean_prompt(class_prompt_collection)[
+            :num_class_images]
+
+    elif concept_type == 'memorization':
+        pipeline.scheduler = DPMSolverMultistepScheduler.from_config(
+            pipeline.scheduler.config)
+        num_prompts_firstpass = 5
+        num_prompts_secondpass = 2
+        threshold = 0.3
+        # Generate num_prompts_firstpass paraphrases which generate different content at least 1-threshold % of the times.
+        os.makedirs(class_images_dir / 'temp/', exist_ok=True)
+        class_prompt_collection_counter = []
+        caption_target = []
+        prev_captions = []
+        messages = [{"role": "user", "content": f"Generate {4*num_prompts_firstpass} different paraphrase of the caption: {class_prompt}. Preserve the meaning when paraphrasing."}]
+        while True:
+            completion = openai.ChatCompletion.create(
+                model="gpt-3.5-turbo",
+                messages=messages
+            )
+            # print(completion.choices[0].message.content.lower().split('\n'))
+            class_prompt_collection_ = [x.strip(
+            ) for x in completion.choices[0].message.content.lower().split('\n') if x.strip() != '']
+            class_prompt_collection_ = clean_prompt(class_prompt_collection_)
+            # print(class_prompt_collection_)
+            for prompt in tqdm(
+                class_prompt_collection_, desc="Generating anchor and target prompts ", disable=not accelerator.is_local_main_process
+            ):
+                print(f'Prompt: {prompt}')
+                images = pipeline([prompt]*10, num_inference_steps=25,).images
+
+                score = filter(images, mem_impath, return_score=True)
+                print(f'Memorization rate: {score}')
+                if score <= threshold and prompt not in class_prompt_collection and len(class_prompt_collection) < num_prompts_firstpass:
+                    class_prompt_collection += [prompt]
+                    class_prompt_collection_counter += [score]
+                elif score >= 0.6 and prompt not in caption_target and len(caption_target) < 2:
+                    caption_target += [prompt]
+                if len(class_prompt_collection) >= num_prompts_firstpass and len(caption_target) >= 2:
+                    break
+
+            if len(class_prompt_collection) >= num_prompts_firstpass:
+                break
+            # print("prompts till now", class_prompt_collection, caption_target)
+            # print("prompts till now", len(
+            #     class_prompt_collection), len(caption_target))
+            prev_captions += class_prompt_collection_
+            prev_captions_ = ','.join(prev_captions[-40:])
+
+            messages = [
+                {"role": "user", "content": f"Generate {4*(num_prompts_firstpass- len(class_prompt_collection))} different paraphrase of the caption: {class_prompt}. Preserve the meaning the most when paraphrasing. Also make sure that the new captions are different from the following captions: {prev_captions_[:4000]}"}]
+
+        # Generate more paraphrases using the captions we retrieved above.
+        for prompt in class_prompt_collection[:num_prompts_firstpass]:
+            completion = openai.ChatCompletion.create(
+                model="gpt-3.5-turbo",
+                messages=[
+                    {"role": "user", "content": f"Generate {num_prompts_secondpass} different paraphrases of: {prompt}. "}]
+
+            )
+            class_prompt_collection += clean_prompt(
+                [x.strip() for x in completion.choices[0].message.content.lower().split('\n') if x.strip() != ''])
+
+        for prompt in tqdm(class_prompt_collection[num_prompts_firstpass:], desc="Memorization rate for final prompts"):
+            images = pipeline([prompt]*10, num_inference_steps=25,).images
+
+            class_prompt_collection_counter += [
+                filter(images, mem_impath, return_score=True)]
+
+        # select least ten and most memorized text prompts to be selected as anchor and target prompts.
+        class_prompt_collection = sorted(
+            zip(class_prompt_collection, class_prompt_collection_counter), key=lambda x: x[1])
+        caption_target += [x for (x, y) in class_prompt_collection if y >= 0.6]
+        class_prompt_collection = [
+            x for (x, y) in class_prompt_collection if y <= threshold][:10]
+        print("Anchor prompts:", class_prompt_collection)
+        print("Target prompts:", caption_target)
+    return class_prompt_collection, ';*+'.join(caption_target)
+
+
+def clean_prompt(class_prompt_collection):
+    class_prompt_collection = [re.sub(
+        r"[0-9]+", lambda num: '' * len(num.group(0)), prompt) for prompt in class_prompt_collection]
+    class_prompt_collection = [re.sub(
+        r"^\.+", lambda dots: '' * len(dots.group(0)), prompt) for prompt in class_prompt_collection]
+    class_prompt_collection = [x.strip() for x in class_prompt_collection]
+    class_prompt_collection = [x.replace('"', '') for x in class_prompt_collection]
+    return class_prompt_collection
+
+
+def safe_dir(dir):
+    if not dir.exists():
+        dir.mkdir()
+    return dir

models/greg_rutkowski_ablation_delta.bin

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

models/vangogh_ablation_delta.bin

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

requirements.txt

@@ -0,0 +1,13 @@
+gradio
+torch
+torchvision
+scipy
+accelerate
+modelcards
+transformers>=4.25.1
+diffusers==0.19.0
+tqdm 
+openai
+triton
+xformers>=0.0.20
+bitsandbytes

trainer.py

@@ -0,0 +1,139 @@
+import gradio as gr
+import PIL.Image
+import shlex
+import shutil
+import subprocess
+from pathlib import Path
+import os
+import torch
+from tqdm import tqdm
+
+
+
+def pad_image(image: PIL.Image.Image) -> PIL.Image.Image:
+    w, h = image.size
+    if w == h:
+        return image
+    elif w > h:
+        new_image = PIL.Image.new(image.mode, (w, w), (0, 0, 0))
+        new_image.paste(image, (0, (w - h) // 2))
+        return new_image
+    else:
+        new_image = PIL.Image.new(image.mode, (h, h), (0, 0, 0))
+        new_image.paste(image, ((h - w) // 2, 0))
+        return new_image
+
+
+
+def train_submit(
+        prompt, anchor_prompt, concept_type, reg_lambda, iterations, lr, openai_key, save_path, mem_impath=None
+    ):
+        if not torch.cuda.is_available():
+            raise gr.Error('CUDA is not available.')
+
+        torch.cuda.empty_cache()
+        original_prompt = prompt
+        parameter_group = "cross-attn"
+        train_batch_size = 4
+        if concept_type == 'style':
+            class_data_dir = f'./data/samples_painting/'
+            anchor_prompt = f'./assets/painting.txt'
+            openai_key = ''
+        elif concept_type == 'object':
+            os.makedirs('temp', exist_ok=True)
+            class_data_dir = f'./temp/{anchor_prompt}'
+            name = save_path.split('/')[-1]
+            prompt = f'{anchor_prompt}+{prompt}'
+            assert openai_key is not None
+
+            if len(openai_key.split('\n')) > 1:
+                openai_key = openai_key.split('\n')
+                with open(f'./temp/{name}.txt', 'w') as f:
+                    for prompt_ in openai_key:
+                        f.write(prompt_.strip()+'\n')
+                openai_key = ''
+                anchor_prompt = f'./temp/{name}.txt'
+        elif concept_type == 'memorization':
+            os.system("wget https://dl.fbaipublicfiles.com/sscd-copy-detection/sscd_imagenet_mixup.torchscript.pt -P assets/") 
+            os.makedirs('temp', exist_ok=True)
+            prompt = f'*+{prompt}'
+            name = save_path.split('/')[-1]
+            train_batch_size = 1
+            lr = 5e-7
+            parameter_group = "full-weight"
+
+            assert openai_key is not None
+            assert mem_impath is not None
+
+            if len(openai_key.split('\n')) > 1:
+                openai_key = openai_key.split('\n')
+                with open(f'./temp/{name}.txt', 'w') as f:
+                    for prompt_ in openai_key:
+                        f.write(prompt_.strip()+'\n')
+                openai_key = ''
+                anchor_prompt = f'./temp/{name}.txt'
+            else:
+                anchor_prompt = prompt
+
+            print(mem_impath)
+            image = PIL.Image.open(mem_impath[0][0].name)
+            image = pad_image(image)
+            image = image.convert('RGB')
+            mem_impath = f"./temp/{original_prompt.lower().replace(' ', '')}.jpg"
+            image.save(mem_impath, format='JPEG', quality=100)
+
+            class_data_dir = f"./temp/{original_prompt.lower().replace(' ', '')}"
+            
+
+        command = f'''
+        accelerate launch concept-ablation-diffusers/train.py \
+          --pretrained_model_name_or_path="CompVis/stable-diffusion-v1-4"  \
+          --output_dir={save_path} \
+          --class_data_dir={class_data_dir} \
+          --class_prompt="{anchor_prompt}"  \
+          --caption_target "{prompt}" \
+          --concept_type {concept_type} \
+          --resolution=512  \
+          --train_batch_size={train_batch_size}  \
+          --learning_rate={lr}  \
+          --max_train_steps={iterations} \
+          --scale_lr --hflip \
+          --parameter_group {parameter_group} \
+          --openai_key "{openai_key}" \
+          --enable_xformers_memory_efficient_attention  --num_class_images 500
+        '''
+
+        if concept_type == 'style':
+            command += f'  --noaug'
+        
+        if concept_type == 'memorization':
+            command += f' --use_8bit_adam  --with_prior_preservation --prior_loss_weight=1.0 --mem_impath {mem_impath}'
+            
+        with open(f'{save_path}/train.sh', 'w') as f:
+            command_s = ' '.join(command.split())
+            f.write(command_s)
+
+        res = subprocess.run(shlex.split(command))
+
+        if res.returncode == 0:
+            result_message = 'Training Completed!'
+        else:
+            result_message = 'Training Failed!'
+        weight_paths = sorted(Path(save_path).glob('*.bin'))
+        print(weight_paths)
+        return gr.update(value=result_message), weight_paths[0]
+
+
+def inference(model_path, prompt, n_steps, generator):
+    import sys 
+    sys.path.append('concept-ablation/diffusers/.')
+    from model_pipeline import CustomDiffusionPipeline
+    import torch
+
+    pipe = CustomDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16).to("cuda")
+    image1 = pipe(prompt, num_inference_steps=n_steps, guidance_scale=6., eta=1., generator=generator).images[0]
+
+    pipe.load_model(model_path)
+    image2 = pipe(prompt, num_inference_steps=n_steps, guidance_scale=6., eta=1., generator=generator).images[0]
+
+    return image1, image2