Update app.py

app.py

@@ -35,13 +35,13 @@ def infer(prompt, init_image):
         text_embedding_optimization_steps=100,
         model_fine_tuning_optimization_steps=500)
     
-    #with torch.no_grad():
-    #    torch.cuda.empty_cache()
+    with torch.no_grad():
+        torch.cuda.empty_cache()
     
-    #res = pipe(alpha=1)
+    res = pipe(alpha=1)
 
-    #return res.images[0]
-    return 'trained success'
+    return res.images[0]
+    #return 'trained success'
 
 title = """
     <div style="text-align: center; max-width: 650px; margin: 0 auto;">
@@ -116,10 +116,10 @@ with gr.Blocks(css=css) as block:
         submit_btn = gr.Button("Train")
         
         image_output = gr.Image(label="Edited image")
-        text_output = gr.Image(label="trained status")
+        #text_output = gr.Image(label="trained status")
         
         gr.HTML(article)
 
-    submit_btn.click(fn=infer, inputs=[prompt_input,image_init], outputs=[text_output])
+    submit_btn.click(fn=infer, inputs=[prompt_input,image_init], outputs=[image_output])
     
 block.queue(max_size=12).launch(show_api=False)