Update app.py

app.py

@@ -277,7 +277,8 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
     generator = torch.Generator(device="cuda").manual_seed(seed)
     with calculateDuration("Generating image"):
         # Generate image
-        for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(prompt=prompt_mash,
+        for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+            prompt=prompt_mash,
             num_inference_steps=steps,
             guidance_scale=cfg_scale,
             width=width,
@@ -288,6 +289,7 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
             good_vae=good_vae,
         ):
             yield img
+    pipe.to("cpu")
 
 @spaces.GPU(duration=70)
 def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps, cfg_scale, width, height, seed):
@@ -306,6 +308,7 @@ def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps
         joint_attention_kwargs={"scale": 1.0},
         output_type="pil",
     ).images[0]
+    pipe_i2i.to("cpu")
     return final_image
 
 @spaces.GPU(duration=70)