Update app.py

app.py

@@ -96,7 +96,7 @@ def sample_then_run(network, unet):
     image = inference( network, unet, prompt, negative_prompt, cfg, steps, seed)
     torch.save(network.proj.detach().cpu(), "model.pt" )
     print("done saving")
-    return image, "model.pt", unet.cpu(), network.cpu()
+    return image, "model.pt", network.proj.detach().cpu()
 
 
     
@@ -479,7 +479,7 @@ with gr.Blocks(css="style.css") as demo:
                     outputs = [input_image, file_output])
         
             
-        sample.click(fn=sample_then_run, outputs=[network, unet, input_image, file_output])
+        sample.click(fn=sample_then_run, outputs=[network, input_image, file_output])
         
         submit.click(
             fn=edit_inference, inputs=[prompt, negative_prompt, cfg, steps, seed, injection_step, a1, a2, a3, a4], outputs=[gallery]