Spaces:
Sleeping
Sleeping
File size: 10,994 Bytes
8173ae1 4b2c52c 8173ae1 4b2c52c 8173ae1 4b2c52c 8173ae1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 |
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=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()
|