Spaces:
Sleeping
Sleeping
| 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', | |
| 'R2D2' : 'models/r2d2_delta.bin', | |
| 'Grumpy Cat' : 'models/grumpy_cat_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 training. 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="a house in the style of van gogh", | |
| 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=200, | |
| 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() | |