Update app.py

app.py

@@ -38,7 +38,7 @@ def train(prompt, init_image, trn_text, trn_steps):
     with torch.no_grad():
         torch.cuda.empty_cache()
     
-    return "Training is finished !"
+    return "Training is finished !", gr.Update(value=0), gr.Update(value=0)
    
 def generate(prompt, init_image, trn_text, trn_steps):
     init_image = Image.open(init_image).convert("RGB")
@@ -151,7 +151,7 @@ with gr.Blocks(css=css) as block:
         
         gr.HTML(article)
 
-    train_btn.click(fn=train, inputs=[prompt_input,image_init,trn_text,trn_steps], outputs=[training_status])
+    train_btn.click(fn=train, inputs=[prompt_input,image_init,trn_text,trn_steps], outputs=[training_status, trn_text, trn_steps])
     gen_btn.click(fn=generate, inputs=[prompt_input,image_init,trn_text,trn_steps], outputs=[image_output])
     
 block.queue(max_size=12).launch(show_api=False)