Update app.py

app.py

@@ -77,6 +77,9 @@ def setup_model(prompt, model, num_iterations, learning_rate, progress=gr.Progre
 
 def generate_image(setup_args, num_iterations):
 
+    """Clear CUDA memory before starting executing task."""
+    torch.cuda.empty_cache()  # Free up cached memory
+
     args = setup_args[0]
     trainer = setup_args[1]
     device = setup_args[2]