Update app.py

app.py

@@ -218,7 +218,7 @@ def sample_then_run():
     steps = 25
     image = inference( prompt, negative_prompt, cfg, steps, seed)
     torch.save(network.value.proj, "model.pt" )
-    return image, "model.pt"
+    return image, "model.pt", network.value
 
 
 class CustomImageDataset(Dataset):
@@ -452,7 +452,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]