Update app.py

app.py

@@ -137,7 +137,7 @@ def sample_then_run(net):
     cfg = 3.0
     steps = 25
     image = inference(net, prompt, negative_prompt, cfg, steps, seed)
-    return net, image
+    return net,net,image
 
 @torch.no_grad()
 @spaces.GPU()
@@ -182,9 +182,9 @@ def inference(net, prompt, negative_prompt, guidance_scale, ddim_steps, seed):
     for i,t in enumerate(tqdm.tqdm(noise_scheduler.timesteps)):
         latent_model_input = torch.cat([latents] * 2)
         latent_model_input = noise_scheduler.scale_model_input(latent_model_input, timestep=t)
-    
-        with network:
-            noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings, timestep_cond= None).sample
+        with torch.no_grad():
+            with network:
+                noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings, timestep_cond= None).sample
             
         #guidance
         noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
@@ -511,7 +511,7 @@ with gr.Blocks(css="style.css") as demo:
         #                 outputs = [input_image, file_output])
             
                 
-        sample.click(fn=sample_then_run,inputs = [net], outputs=[net, input_image])
+        sample.click(fn=sample_then_run,inputs = [net], outputs=[net, file_output, input_image])
             
         submit.click(
                 fn=edit_inference, inputs=[net, prompt, negative_prompt, cfg, steps, seed, injection_step, a1, a2, a3, a4], outputs=[net, gallery]