app.py

import gradio as gr
import torch
import os

from diffusers.utils import is_xformers_available

from finetuning import FineTunedModel
from StableDiffuser import StableDiffuser
from memory_efficiency import MemoryEfficiencyWrapper
from train import train, training_should_cancel
    
import os

model_map = {}
model_names_list = []

def populate_global_model_map():
    global model_map
    global model_names_list
    for model_file in os.listdir('models'):
        path = 'models/' + model_file
        if any([existing_path == path for existing_path in model_map.values()]):
            continue
        model_map[model_file] = path
    model_names_list.clear()
    model_names_list.extend(model_map.keys())

populate_global_model_map()

ORIGINAL_SPACE_ID = 'baulab/Erasing-Concepts-In-Diffusion'
SPACE_ID = os.getenv('SPACE_ID')

SHARED_UI_WARNING = f'''## Attention - Training using the ESD-u method does not work in this shared UI. You can either duplicate and use it with a gpu with at least 40GB, or clone this repository to run on your own machine.
<center><a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true"><img style="margin-top:0;margin-bottom:0" src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></center>
'''

# work around Gradio's weird threading

class Demo:

    def __init__(self) -> None:

        self.training = False
        self.generating = False

        with gr.Blocks() as demo:
            self.layout()
            demo.queue(concurrency_count=5).launch()


    def layout(self):

        with gr.Row():

            if SPACE_ID == ORIGINAL_SPACE_ID:

                self.warning = gr.Markdown(SHARED_UI_WARNING)
          
        with gr.Row():
                
            with gr.Tab("Test") as inference_column:

                with gr.Row():

                    self.explain_infr = gr.Markdown(interactive=False, 
                                      value='This is a demo of [Erasing Concepts from Stable Diffusion](https://erasing.baulab.info/).  To try out a model where a concept has been erased, select a model and enter any prompt.  For example, if you select the model "Van Gogh" you can generate images for the prompt "A portrait in the style of Van Gogh" and compare the erased and unerased models.  We have also provided several other pre-fine-tuned models with artistic styles and objects erased (Check out the "ESD Model" drop-down). You can also train and run your own custom models. Check out the "train" section for custom erasure of concepts.')

                with gr.Row():

                    with gr.Column(scale=1):

                        self.base_repo_id_or_path_input_infr = gr.Text(
                            label="Base model",
                            value="CompVis/stable-diffusion-v1-4",
                            info="Path or huggingface repo id of the base model that this edit was done against"
                        )

                        self.prompt_input_infr = gr.Text(
                            placeholder="Enter prompt...",
                            label="Prompt",
                            info="Prompt to generate"
                        )
                        self.negative_prompt_input_infr = gr.Text(
                            label="Negative prompt"
                        )
                        self.seed_infr = gr.Number(
                            label="Seed",
                            value=42
                        )
                        with gr.Row():
                            self.img_width_infr = gr.Slider(
                                label="Image width",
                                minimum=256,
                                maximum=1024,
                                value=512,
                                step=64
                            )
                            self.img_height_infr = gr.Slider(
                                label="Image height",
                                minimum=256,
                                maximum=1024,
                                value=512,
                                step=64
                            )

                        with gr.Row():
                            self.model_dropdown = gr.Dropdown(
                                label="ESD Model",
                                choices= list(model_map.keys()),
                                value='Van Gogh',
                                interactive=True
                            )
                            self.model_reload_button = gr.Button(
                                value="🔄",
                                interactive=True
                            )

                    with gr.Column(scale=2):

                        self.infr_button = gr.Button(
                            value="Generate",
                            interactive=True
                        )

                        with gr.Row():

                            self.image_new = gr.Image(
                                label="ESD",
                                interactive=False
                            )
                            self.image_orig = gr.Image(
                                label="SD",
                                interactive=False
                            )

            with gr.Tab("Train") as training_column:

                with gr.Row():
                    self.explain_train= gr.Markdown(interactive=False, 
                                      value='In this part you can erase any concept from Stable Diffusion.   Enter a prompt for the concept or style you want to erase, and select ESD-x if you want to focus erasure on prompts that mention the concept explicitly. [NOTE: ESD-u is currently unavailable in this space. But you can duplicate the space and run it on GPU with VRAM >40GB for enabling ESD-u]. With default settings, it takes about 15 minutes to fine-tune the model; then you can try inference above or download the weights.  The training code used here is slightly different than the code tested in the original paper.  Code and details are at [github link](https://github.com/rohitgandikota/erasing).')

                with gr.Row():

                    with gr.Column(scale=3):
                        self.train_model_input = gr.Text(
                            label="Model to Edit",
                            value="CompVis/stable-diffusion-v1-4",
                            info="Path or huggingface repo id of the model to edit"
                        )

                        self.train_img_size_input = gr.Slider(
                            value=512,
                            step=64,
                            minimum=256,
                            maximum=1024,
                            label="Image Size",
                            info="Image size for training, should match the model's native image size"
                        )

                        self.train_prompts_input = gr.Text(
                            placeholder="Enter prompts, one per line",
                            label="Prompts to Erase",
                            info="Prompts corresponding to concepts to erase, one per line"
                        )

                        choices = ['ESD-x', 'ESD-self', 'ESD-u']
                        #if torch.cuda.get_device_properties(0).total_memory * 1e-9 >= 40 or is_xformers_available():
                        #    choices.append('ESD-u')
                    
                        self.train_method_input = gr.Dropdown(
                            choices=choices,
                            value='ESD-x',
                            label='Train Method',
                            info='Method of training. ESD-x uses the least VRAM, and you may get OOM errors with the other methods.'
                        )

                        self.neg_guidance_input = gr.Number(
                            value=1,
                            label="Negative Guidance",
                            info='Guidance of negative training used to train'
                        )

                        self.iterations_input = gr.Number(
                            value=150,
                            precision=0,
                            label="Iterations",
                            info='iterations used to train'
                        )

                        self.lr_input = gr.Number(
                            value=1e-5,
                            label="Learning Rate",
                            info='Learning rate used to train'
                        )
                        self.train_seed_input = gr.Number(
                            value=-1,
                            label="Seed",
                            info="Set to a fixed number for reproducible training results, or use -1 to pick randomly"
                        )
                        self.train_save_every_input = gr.Number(
                            value=-1,
                            label="Save Every N Steps",
                            info="If >0, save the model throughout training at the given step interval."
                        )

                        with gr.Column():
                            self.train_memory_options = gr.Markdown(interactive=False,
                                value='Performance and VRAM usage optimizations, may not work on all devices:')
                            with gr.Row():
                                self.train_use_adamw8bit_input = gr.Checkbox(label="8bit AdamW", value=True)
                                self.train_use_xformers_input = gr.Checkbox(label="xformers", value=True)
                                self.train_use_amp_input = gr.Checkbox(label="AMP", value=True)
                                self.train_use_gradient_checkpointing_input = gr.Checkbox(
                                    label="Gradient checkpointing", value=False)

                        self.train_validation_prompts = gr.TextArea(
                            label="Validation Prompts",
                            placeholder="Probably, you want to put the \"Prompt to Erase\" in here as the first entry...",
                            value='',
                            info="Prompts for producing validation graphs, one per line."
                        )
                        self.train_sample_positive_prompts = gr.TextArea(
                            label="Sample Prompts",
                            value='',
                            info="Positive prompts for generating sample images, one per line."
                        )
                        self.train_sample_negative_prompts = gr.TextArea(
                            label="Sample Negative Prompts",
                            value='',
                            info="Negative prompts for use when generating sample images. One for each positive prompt, or leave empty for none."
                        )

                        with gr.Row():
                            self.train_sample_batch_size_input = gr.Slider(
                                value=1,
                                step=1,
                                minimum=1,
                                maximum=32,
                                label="Sample generation batch size",
                                info="Batch size for sample generation, larger needs more VRAM"
                            )
                            self.train_validate_every_n_steps = gr.Number(
                                label="Validate Every N Steps",
                                value=20,
                                info="Validation and sample generation will be run at intervals of this many steps"
                            )

                    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.train_cancel_button = gr.Button(
                            value="Cancel Training"
                        )

                        self.download = gr.Files()

            with gr.Tab("Export") as export_column:
                with gr.Row():
                    self.explain_train= gr.Markdown(interactive=False,
                        value='Export a model to Diffusers format. Please enter the base model and select the editing weights.')

                with gr.Row():

                    with gr.Column(scale=3):
                        self.base_repo_id_or_path_input_export = gr.Text(
                            label="Base model",
                            value="CompVis/stable-diffusion-v1-4",
                            info="Path or huggingface repo id of the base model that this edit was done against"
                        )

                        with gr.Row():
                            self.model_dropdown_export = gr.Dropdown(
                                label="ESD Model",
                                choices=list(model_map.keys()),
                                value='Van Gogh',
                                interactive=True
                            )
                            self.model_reload_button_export = gr.Button(
                                value="🔄",
                                interactive=True
                            )

                        self.save_path_input_export = gr.Text(
                            label="Output path",
                            placeholder="./exported_models/model_name",
                            info="Path to export the model to. A diffusers folder will be written to this location."
                        )

                        self.save_half_export = gr.Checkbox(
                            label="Save as fp16"
                        )

                    with gr.Column(scale=1):
                        self.export_status = gr.Button(
                            value='', variant='primary', label='Status', interactive=False)
                        self.export_button = gr.Button(
                            value="Export")
                        self.export_download = gr.Files()

        self.infr_button.click(self.inference, inputs = [
            self.prompt_input_infr,
            self.negative_prompt_input_infr,
            self.seed_infr,
            self.img_width_infr,
            self.img_height_infr,
            self.model_dropdown,
            self.base_repo_id_or_path_input_infr
            ],
            outputs=[
                self.image_new,
                self.image_orig
            ]
        )
        self.model_reload_button.click(self.reload_models,
                                       inputs=[self.model_dropdown, self.model_dropdown_export],
                                       outputs=[self.model_dropdown, self.model_dropdown_export])

        self.model_reload_button_export.click(self.reload_models,
                                       inputs=[self.model_dropdown, self.model_dropdown_export],
                                       outputs=[self.model_dropdown, self.model_dropdown_export])
        train_event = self.train_button.click(self.train, inputs = [
            self.train_model_input,
            self.train_img_size_input,
            self.train_prompts_input,
            self.train_method_input, 
            self.neg_guidance_input,
            self.iterations_input,
            self.lr_input,
            self.train_use_adamw8bit_input,
            self.train_use_xformers_input,
            self.train_use_amp_input,
            self.train_use_gradient_checkpointing_input,
            self.train_seed_input,
            self.train_save_every_input,
            self.train_sample_batch_size_input,
            self.train_validation_prompts,
            self.train_sample_positive_prompts,
            self.train_sample_negative_prompts,
            self.train_validate_every_n_steps
        ],
        outputs=[self.train_button, self.train_status, self.download, self.model_dropdown]
        )
        self.train_cancel_button.click(self.cancel_training,
                                       inputs=[],
                                       outputs=[self.train_cancel_button],
                                       cancels=[train_event])

        self.export_button.click(self.export, inputs = [
            self.model_dropdown_export,
            self.base_repo_id_or_path_input_export,
            self.save_path_input_export,
            self.save_half_export
        ],
        outputs=[self.export_button, self.export_status]
        )

    def reload_models(self, model_dropdown, model_dropdown_export):
        current_model_name = model_dropdown
        current_model_name_export = model_dropdown_export
        populate_global_model_map()
        global model_names_list
        return [gr.update(choices=model_names_list, value=current_model_name),
                gr.update(choices=model_names_list, value=current_model_name_export)]

    def cancel_training(self):
        if self.training:
            training_should_cancel.release()
            print("cancellation requested...")
        return [gr.update(value="Cancelling...", interactive=True)]

    def train(self, repo_id_or_path, img_size, prompts, train_method, neg_guidance, iterations, lr,
              use_adamw8bit=True, use_xformers=False, use_amp=False, use_gradient_checkpointing=False,
              seed=-1, save_every=-1, sample_batch_size=1,
              validation_prompts: str=None, sample_positive_prompts: str=None, sample_negative_prompts: str=None, validate_every_n_steps=-1,
              pbar=gr.Progress(track_tqdm=True)):
        """

        :param repo_id_or_path:
        :param img_size:
        :param prompts:
        :param train_method:
        :param neg_guidance:
        :param iterations:
        :param lr:
        :param use_adamw8bit:
        :param use_xformers:
        :param use_amp:
        :param use_gradient_checkpointing:
        :param seed:
        :param save_every:
        :param validation_prompts: split on \n
        :param sample_positive_prompts: split on \n
        :param sample_negative_prompts: split on \n
        :param validate_every_n_steps: split on \n
        :param pbar:
        :return:
        """
        if self.training:
            return [gr.update(interactive=True, value='Train'), gr.update(value='Someone else is training... Try again soon'), None, gr.update()]

        print(f"Training {repo_id_or_path} at {img_size} to remove '{prompts}'.")
        print(f"  {train_method}, negative guidance {neg_guidance}, lr {lr}, {iterations} iterations.")
        print(f" {'✅' if use_gradient_checkpointing else '❌'} gradient checkpointing")
        print(f" {'✅' if use_amp else '❌'} AMP")
        print(f" {'✅' if use_xformers else '❌'} xformers")
        print(f" {'✅' if use_adamw8bit else '❌'} 8-bit AdamW")

        if train_method == 'ESD-x':
            modules = ".*attn2$"
            frozen = []

        elif train_method == 'ESD-u':
            modules = "unet$"
            frozen = [".*attn2$", "unet.time_embedding$", "unet.conv_out$"]   

        elif train_method == 'ESD-self':
            modules = ".*attn1$"
            frozen = []

        # build a save path, ensure it isn't in use
        while True:
            randn = torch.randint(1, 10000000, (1,)).item()
            options = f'{"a8" if use_adamw8bit else ""}{"AM" if use_amp else ""}{"xf" if use_xformers else ""}{"gc" if use_gradient_checkpointing else ""}'
            save_path = f"models/{prompts[0].lower().replace(' ', '')}_{train_method}_ng{neg_guidance}_lr{lr}_iter{iterations}_seed{seed}_{options}__{randn}.pt"
            if not os.path.exists(save_path):
                break
            # repeat until a not-in-use path is found

        prompts = [p for p in prompts.split('\n') if len(p)>0]
        validation_prompts = [] if validation_prompts is None else [p for p in validation_prompts.split('\n') if len(p)>0]
        sample_positive_prompts = [] if sample_positive_prompts is None else [p for p in sample_positive_prompts.split('\n') if len(p)>0]
        sample_negative_prompts = [] if sample_negative_prompts is None else sample_negative_prompts.split('\n')
        print(f"validation prompts: {validation_prompts}")
        print(f"sample positive prompts: {sample_positive_prompts}")
        print(f"sample negative prompts: {sample_negative_prompts}")

        try:
            self.training = True
            self.train_cancel_button.update(interactive=True)
            batch_size = 1 # other batch sizes are non-functional
            save_path = train(repo_id_or_path, img_size, prompts, modules, frozen, iterations, neg_guidance, lr, save_path,
                  use_adamw8bit, use_xformers, use_amp, use_gradient_checkpointing,
                  seed=int(seed), save_every_n_steps=int(save_every),
                              batch_size=int(batch_size), sample_batch_size=int(sample_batch_size),
                  validate_every_n_steps=validate_every_n_steps, validation_prompts=validation_prompts,
                  sample_positive_prompts=sample_positive_prompts, sample_negative_prompts=sample_negative_prompts)
            if save_path is None:
                new_model_name = None
                finished_message = "Training cancelled."
            else:
                new_model_name = f'{os.path.basename(save_path)}'
                finished_message = f'Done Training! Try your model ({new_model_name}) in the "Test" tab'
        finally:
            self.training = False
            self.train_cancel_button.update(interactive=False)

        torch.cuda.empty_cache()

        if new_model_name is not None:
            model_map[new_model_name] = save_path

        return [gr.update(interactive=True, value='Train'),
                gr.update(value=finished_message),
                save_path,
                gr.Dropdown.update(choices=list(model_map.keys()), value=new_model_name)]

    def export(self, model_name, base_repo_id_or_path, save_path, save_half):
        model_path = model_map[model_name]
        checkpoint = torch.load(model_path)
        diffuser = StableDiffuser(scheduler='DDIM',
                                       keep_pipeline=True,
                                       repo_id_or_path=base_repo_id_or_path,
                                       ).eval()
        finetuner = FineTunedModel.from_checkpoint(diffuser, checkpoint).eval()
        with finetuner:
            if save_half:
                diffuser = diffuser.half()
                diffuser.pipeline.to('cpu', torch_dtype=torch.float16)
            diffuser.pipeline.save_pretrained(save_path)

        return [gr.update(interactive=True, value='Export'),
                gr.update(value=f'Done Exporting! Diffusers folder is at {os.path.realpath(save_path)}.')]


    def inference(self, prompt, negative_prompt, seed, width, height, model_name, base_repo_id_or_path, pbar = gr.Progress(track_tqdm=True)):
        
        seed = seed or 42
        model_path = model_map[model_name]
        checkpoint = torch.load(model_path)

        if type(prompt) is str:
            prompt = [prompt]
        if type(negative_prompt) is str:
            negative_prompt = [negative_prompt]

        self.diffuser = StableDiffuser(scheduler='DDIM', repo_id_or_path=base_repo_id_or_path).to('cuda').eval().half()
        finetuner = FineTunedModel.from_checkpoint(self.diffuser, checkpoint).eval().half()

        generator = torch.manual_seed(seed)

        torch.cuda.empty_cache()
        images = self.diffuser(
            prompt,
            negative_prompt,
            width=width,
            height=height,
            n_steps=50,
            generator=generator
        )
        orig_image = images[0][0]

        torch.cuda.empty_cache()
        with finetuner:
            images = self.diffuser(
                prompt,
                negative_prompt,
                width=width,
                height=height,
                n_steps=50,
                generator=generator
            )
        edited_image = images[0][0]

        del finetuner
        torch.cuda.empty_cache()

        return edited_image, orig_image


demo = Demo()