amildravid4292 commited on
Commit
2d4497e
·
verified ·
1 Parent(s): c920a8b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -1
app.py CHANGED
@@ -94,9 +94,14 @@ def sample_then_run(network, unet):
94
  cfg = 3.0
95
  steps = 25
96
  image = inference( network, unet, prompt, negative_prompt, cfg, steps, seed)
97
- torch.save(network.proj.cpu().detach(), "model.pt" )
 
98
  return image, "model.pt", unet, network
99
 
 
 
 
 
100
 
101
 
102
  @torch.no_grad()
 
94
  cfg = 3.0
95
  steps = 25
96
  image = inference( network, unet, prompt, negative_prompt, cfg, steps, seed)
97
+ save_model(network)
98
+ #torch.save(network.proj.cpu().detach(), "model.pt" )
99
  return image, "model.pt", unet, network
100
 
101
+ def save_model(network):
102
+ torch.save(network.proj.cpu().detach(), "model.pt" )
103
+
104
+
105
 
106
 
107
  @torch.no_grad()