Merge branch 'main' of https://huggingface.co/spaces/baulab/Erasing-Concepts-In-Diffusion into main

app.py

@@ -258,7 +258,8 @@ class Demo:
             loss.backward()
             optimizer.step()
 
-        torch.save(finetuner.state_dict(), 'ft.pt')
+        ft_path = f"{prompt.lower().replace(' ', '')}.pt"
+        torch.save(finetuner.state_dict(), ft_path)
 
         self.finetuner = finetuner.eval().half()
 
@@ -268,9 +269,9 @@ class Demo:
 
         self.training = False      
 
-        model_map['Custom'] = 'ft.pt'      
+        model_map['Custom'] = ft_path
 
-        return [gr.update(interactive=True), gr.update(interactive=True), 'ft.pt', gr.Dropdown.update(choices=list(model_map.keys()), value='Custom')]
+        return [gr.update(interactive=True), gr.update(interactive=True), ft_path, gr.Dropdown.update(choices=list(model_map.keys()), value='Custom')]
 
 
     def inference(self, prompt, seed, pbar = gr.Progress(track_tqdm=True)):