working?

app.py

@@ -7,104 +7,110 @@ from omegaconf import OmegaConf
 from StableDiffuser import StableDiffuser
 from diffusers import UNet2DConditionModel
 
-ckpt_path = "stable-diffusion/models/ldm/sd-v1-4-full-ema.ckpt"
-config_path = "stable-diffusion/configs/stable-diffusion/v1-inference.yaml"
-diffusers_config_path = "stable-diffusion/config.json"
+ckpt_path = "stable_diffusion/models/ldm/sd-v1-4-full-ema.ckpt"
+config_path = "stable_diffusion/configs/stable-diffusion/v1-inference.yaml"
+diffusers_config_path = "stable_diffusion/config.json"
 
 
 class Demo:
 
     def __init__(self) -> None:
-        demo = self.layout()
-        demo.launch()
 
+        with gr.Blocks() as demo:
+            self.layout()
+            demo.queue(concurrency_count=10).launch()
+
+    def disable(self):
+        return [gr.update(interactive=False), gr.update(interactive=False)]
 
     def layout(self):
+        with gr.Row():
+            with gr.Column() as training_column:
+                self.prompt_input = gr.Text(
+                    placeholder="Enter prompt...",
+                    label="Prompt",
+                    info="Prompt corresponding to concept to erase"
+                )
+                self.train_method_input = gr.Dropdown(
+                    choices=['noxattn', 'selfattn', 'xattn', 'full'],
+                    value='xattn',
+                    label='Train Method',
+                    info='Method of training'
+                )
 
-        with gr.Blocks() as demo:
+                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=1000,
+                    precision=0,
+                    label="Iterations",
+                    info='iterations used to train'
+                )
 
-            with gr.Row():
-                with gr.Column() as training_column:
-                    self.prompt_input = gr.Text(
+                self.lr_input = gr.Number(
+                    value=1e-5,
+                    label="Learning Rate",
+                    info='Learning rate used to train'
+                )
+                self.train_button = gr.Button(
+                    value="Train",
+                    
+                )
+                
+                
+            with gr.Column() as inference_column:
+
+                with gr.Row():
+
+                    self.prompt_input_infr = gr.Text(
                         placeholder="Enter prompt...",
                         label="Prompt",
                         info="Prompt corresponding to concept to erase"
                     )
-                    self.train_method_input = gr.Dropdown(
-                        choices=['noxattn', 'selfattn', 'xattn', 'full'],
-                        value='xattn',
-                        label='Train Method',
-                        info='Method of training'
-                    )
 
-                    self.neg_guidance_input = gr.Number(
-                        value=1,
-                        label="Negative Guidance",
-                        info='Guidance of negative training used to train'
-                    )
+                with gr.Row():
 
-                    self.iterations_input = gr.Number(
-                        value=1000,
-                        precision=0,
-                        label="Iterations",
-                        info='iterations used to train'
+                    self.image_new = gr.Image(
+                        label="New Image",
+                        interactive=False
                     )
-
-                    self.lr_input = gr.Number(
-                        value=1e-5,
-                        label="Learning Rate",
-                        info='Learning rate used to train'
+                    self.image_orig = gr.Image(
+                        label="Orig Image",
+                        interactive=False
                     )
 
-                    self.train_button = gr.Button(
-                        value="Train",
+                with gr.Row():
+
+                    self.infr_button = gr.Button(
+                        value="Generate",
+                        interactive=False
                     )
-                    self.train_button.click(self.train, inputs = [
-                            self.prompt_input,
-                            self.train_method_input, 
-                            self.neg_guidance_input,
-                            self.iterations_input,
-                            self.lr_input
+                    self.infr_button.click(self.inference, inputs = [
+                            self.prompt_input_infr,
+                        ],
+                        outputs=[
+                            self.image_new,
+                            self.image_orig
                         ]
                     )
-                with gr.Column() as inference_column:
-
-                    with gr.Row():
-
-                        self.prompt_input_infr = gr.Text(
-                            placeholder="Enter prompt...",
-                            label="Prompt",
-                            info="Prompt corresponding to concept to erase"
-                        )
-
-                    with gr.Row():
-
-                        self.image_new = gr.Image(
-                            label="New Image",
-                            interactive=False
-                        )
-                        self.image_orig = gr.Image(
-                            label="Orig Image",
-                            interactive=False
-                        )
-
-                    with gr.Row():
-
-                        self.infr_button = gr.Button(
-                            value="Generate",
-                        )
-                        self.infr_button.click(self.inference, inputs = [
-                                self.prompt_input_infr,
-                            ],
-                            outputs=[
-                                self.image_new,
-                                self.image_orig
-                            ]
-                        )
-        return demo
-
+                    self.train_button.click(self.disable, 
+                        outputs=[self.train_button, self.infr_button]
+                    )
+                    self.train_button.click(self.train, inputs = [
+                        self.prompt_input,
+                        self.train_method_input, 
+                        self.neg_guidance_input,
+                        self.iterations_input,
+                        self.lr_input
+                    ],
+                    outputs=[self.train_button, self.infr_button]
+                )
 
-    def train(self, prompt, train_method, neg_guidance, iterations, lr):
+    def train(self, prompt, train_method, neg_guidance, iterations, lr, pbar = gr.Progress(track_tqdm=True)):
 
         model_orig, model_edited = train_esd(prompt,
                 train_method,
@@ -115,8 +121,7 @@ class Demo:
                 config_path,
                 ckpt_path, 
                 diffusers_config_path,
-                ['cuda', 'cuda'],
-                gr.Progress()
+                ['cuda', 'cuda']
                 )
         
         original_config = OmegaConf.load(config_path)
@@ -127,6 +132,8 @@ class Demo:
 
         self.init_inference(model_edited_sd, model_orig_sd, unet_config)
 
+        return [gr.update(interactive=True), gr.update(interactive=True)]
+
     def init_inference(self, model_edited_sd, model_orig_sd, unet_config):
 
         self.model_edited_sd = model_edited_sd
@@ -163,5 +170,5 @@ class Demo:
         return edited_image, orig_image
 
 
-
+demo = Demo()
 

requirements.txt

@@ -7,5 +7,7 @@ transformers
 pytorch_lightning==1.6.5
 taming-transformers
 kornia
+scipy
+accelerate
 git+https://github.com/openai/CLIP.git@main#egg=clip
 git+https://github.com/davidbau/baukit.git

train_esd.py

@@ -102,7 +102,7 @@ def get_models(config_path, ckpt_path, devices):
 
     return model_orig, sampler_orig, model, sampler
 
-def train_esd(prompt, train_method, start_guidance, negative_guidance, iterations, lr, config_path, ckpt_path, diffusers_config_path, devices, progress_bar, seperator=None, image_size=512, ddim_steps=50):
+def train_esd(prompt, train_method, start_guidance, negative_guidance, iterations, lr, config_path, ckpt_path, diffusers_config_path, devices, seperator=None, image_size=512, ddim_steps=50):
     '''
     Function to train diffusion models to erase concepts from model weights