app.py

import gradio as gr
from train_esd import train_esd


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"

def train(prompt, train_method, neg_guidance, iterations, lr):

 train_esd(prompt,
 train_method,
 3,
 neg_guidance,
 iterations,
 lr,
 config_path,
 ckpt_path, 
 diffusers_config_path,
 ['cuda']
 )


with gr.Blocks() as demo:

 prompt_input = gr.Text(
 placeholder="Enter prompt...",
 label="Prompt",
 info="Prompt corresponding to concept to erase"
 )
 train_method_input = gr.Dropdown(
 choices=['noxattn', 'selfattn', 'xattn', 'full'],
 value='xattn',
 label='Train Method',
 info='Method of training'
 )

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

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

 lr_input = gr.Number(
 value=1e-5,
 label="Iterations",
 info='Learning rate used to train'
 )

 train_button = gr.Button(
 value="Train",
 )
 train_button.click(train, inputs = [
 prompt_input,
 train_method_input, 
 neg_guidance_input,
 iterations_input,
 lr_input
 ]
 )




demo.launch()