Update app.py

app.py

@@ -304,8 +304,9 @@ def sample_then_run():
     steps = 25
     image = inference( prompt, negative_prompt, cfg, steps, seed)
     torch.save(network.value.proj, "model.pt" )
+    net = network.value.cpu()
     
-    return image, "model.pt"   #, network.value.cpu()
+    return image, "model.pt", net   #, network.value.cpu()
 
 
 class CustomImageDataset(Dataset):
@@ -539,7 +540,7 @@ with gr.Blocks(css="style.css") as demo:
                     outputs = [input_image, file_output])
         
             
-        sample.click(fn=sample_then_run, outputs=[input_image, file_output])
+        sample.click(fn=sample_then_run, outputs=[input_image, file_output, network])
         
         submit.click(
             fn=edit_inference, inputs=[prompt, negative_prompt, cfg, steps, seed, injection_step, a1, a2, a3, a4], outputs=[gallery]