baulab/Erasing-Concepts-In-Diffusion
49
1import gradio as gr2import torch3from finetuning import FineTunedModel4from StableDiffuser import StableDiffuser5from train import train6 7import os8model_map = {'Van Gogh' : 'models/vangogh.pt', 9 'Pablo Picasso': 'models/pablopicasso.pt',10 'Car' : 'models/car.pt',11 'Garbage Truck': 'models/garbagetruck.pt',12 'French Horn': 'models/frenchhorn.pt',13 'Kilian Eng' : 'models/kilianeng.pt',14 'Thomas Kinkade' : 'models/thomaskinkade.pt',15 'Tyler Edlin' : 'models/tyleredlin.pt',16 'Kelly McKernan': 'models/kellymckernan.pt',17 'Rembrandt': 'models/rembrandt.pt' }18 19ORIGINAL_SPACE_ID = 'baulab/Erasing-Concepts-In-Diffusion'20SPACE_ID = os.getenv('SPACE_ID')21 22SHARED_UI_WARNING = f'''## Attention - Training using the ESD-u method does not work in this shared UI. You can either duplicate and use it with a gpu with at least 40GB, or clone this repository to run on your own machine.23<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>24'''25 26 27class Demo:28 29 def __init__(self) -> None:30 31 self.training = False32 self.generating = False33 34 self.diffuser = StableDiffuser(scheduler='DDIM').to('cuda').eval().half()35 36 with gr.Blocks() as demo:37 self.layout()38 demo.queue(concurrency_count=5).launch()39 40 41 def layout(self):42 43 with gr.Row():44 45 if SPACE_ID == ORIGINAL_SPACE_ID:46 47 self.warning = gr.Markdown(SHARED_UI_WARNING)48 49 with gr.Row():50 51 with gr.Tab("Test") as inference_column:52 53 with gr.Row():54 55 self.explain_infr = gr.Markdown(interactive=False, 56 value='This is a demo of [Erasing Concepts from Stable Diffusion](https://erasing.baulab.info/). To try out a model where a concept has been erased, select a model and enter any prompt. For example, if you select the model "Van Gogh" you can generate images for the prompt "A portrait in the style of Van Gogh" and compare the erased and unerased models. We have also provided several other pre-fine-tuned models with artistic styles and objects erased (Check out the "ESD Model" drop-down). You can also train and run your own custom models. Check out the "train" section for custom erasure of concepts.')57 58 with gr.Row():59 60 with gr.Column(scale=1):61 62 self.prompt_input_infr = gr.Text(63 placeholder="Enter prompt...",64 label="Prompt",65 info="Prompt to generate"66 )67 68 with gr.Row():69 70 self.model_dropdown = gr.Dropdown(71 label="ESD Model",72 choices= list(model_map.keys()),73 value='Van Gogh',74 interactive=True75 )76 77 self.seed_infr = gr.Number(78 label="Seed",79 value=4280 )81 82 with gr.Column(scale=2):83 84 self.infr_button = gr.Button(85 value="Generate",86 interactive=True87 )88 89 with gr.Row():90 91 self.image_new = gr.Image(92 label="ESD",93 interactive=False94 )95 self.image_orig = gr.Image(96 label="SD",97 interactive=False98 )99 100 with gr.Tab("Train") as training_column:101 102 with gr.Row():103 104 self.explain_train= gr.Markdown(interactive=False, 105 value='In this part you can erase any concept from Stable Diffusion. Enter a prompt for the concept or style you want to erase, and select ESD-x if you want to focus erasure on prompts that mention the concept explicitly. [NOTE: ESD-u is currently unavailable in this space. But you can duplicate the space and run it on GPU with VRAM >40GB for enabling ESD-u]. With default settings, it takes about 15 minutes to fine-tune the model; then you can try inference above or download the weights. The training code used here is slightly different than the code tested in the original paper. Code and details are at [github link](https://github.com/rohitgandikota/erasing).')106 107 with gr.Row():108 109 with gr.Column(scale=3):110 111 self.prompt_input = gr.Text(112 placeholder="Enter prompt...",113 label="Prompt to Erase",114 info="Prompt corresponding to concept to erase"115 )116 117 choices = ['ESD-x']118 if torch.cuda.get_device_properties(0).total_memory * 1e-9 >= 40:119 choices.append('ESD-u')120 121 self.train_method_input = gr.Dropdown(122 choices=choices,123 value='ESD-x',124 label='Train Method',125 info='Method of training'126 )127 128 self.neg_guidance_input = gr.Number(129 value=1,130 label="Negative Guidance",131 info='Guidance of negative training used to train'132 )133 134 self.iterations_input = gr.Number(135 value=150,136 precision=0,137 label="Iterations",138 info='iterations used to train'139 )140 141 self.lr_input = gr.Number(142 value=1e-5,143 label="Learning Rate",144 info='Learning rate used to train'145 )146 147 with gr.Column(scale=1):148 149 self.train_status = gr.Button(value='', variant='primary', label='Status', interactive=False)150 151 self.train_button = gr.Button(152 value="Train",153 )154 155 self.download = gr.Files()156 157 self.infr_button.click(self.inference, inputs = [158 self.prompt_input_infr,159 self.seed_infr,160 self.model_dropdown161 ],162 outputs=[163 self.image_new,164 self.image_orig165 ]166 )167 self.train_button.click(self.train, inputs = [168 self.prompt_input,169 self.train_method_input, 170 self.neg_guidance_input,171 self.iterations_input,172 self.lr_input173 ],174 outputs=[self.train_button, self.train_status, self.download, self.model_dropdown]175 )176 177 def train(self, prompt, train_method, neg_guidance, iterations, lr, pbar = gr.Progress(track_tqdm=True)):178 179 if self.training:180 return [gr.update(interactive=True, value='Train'), gr.update(value='Someone else is training... Try again soon'), None, gr.update()]181 182 if train_method == 'ESD-x':183 184 modules = ".*attn2$"185 frozen = []186 187 elif train_method == 'ESD-u':188 189 modules = "unet$"190 frozen = [".*attn2$", "unet.time_embedding$", "unet.conv_out$"] 191 192 elif train_method == 'ESD-self':193 194 modules = ".*attn1$"195 frozen = []196 197 randn = torch.randint(1, 10000000, (1,)).item()198 199 save_path = f"models/{randn}_{prompt.lower().replace(' ', '')}.pt"200 201 self.training = True202 203 train(prompt, modules, frozen, iterations, neg_guidance, lr, save_path)204 205 self.training = False206 207 torch.cuda.empty_cache()208 209 model_map['Custom'] = save_path210 211 return [gr.update(interactive=True, value='Train'), gr.update(value='Done Training! \n Try your custom model in the "Test" tab'), save_path, gr.Dropdown.update(choices=list(model_map.keys()), value='Custom')]212 213 214 def inference(self, prompt, seed, model_name, pbar = gr.Progress(track_tqdm=True)):215 216 seed = seed or 42217 218 generator = torch.manual_seed(seed)219 220 model_path = model_map[model_name]221 222 checkpoint = torch.load(model_path)223 224 finetuner = FineTunedModel.from_checkpoint(self.diffuser, checkpoint).eval().half()225 226 torch.cuda.empty_cache()227 228 images = self.diffuser(229 prompt,230 n_steps=50,231 generator=generator232 )233 234 235 orig_image = images[0][0]236 237 torch.cuda.empty_cache()238 239 generator = torch.manual_seed(seed)240 241 with finetuner:242 243 images = self.diffuser(244 prompt,245 n_steps=50,246 generator=generator247 )248 249 edited_image = images[0][0]250 251 del finetuner252 torch.cuda.empty_cache()253 254 return edited_image, orig_image255 256 257demo = Demo()258 259