spacesgpu

app.py

@@ -110,6 +110,10 @@ def generate_flux_image(
         joint_attention_kwargs={"scale": lora_scale},
     ).images[0]
     
+    # Move model to CPU after generation
+    flux_pipeline.to("cpu")
+    torch.cuda.empty_cache()
+    
     # Save the generated image
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     unique_id = str(uuid.uuid4())[:8]