Update app.py

app.py

@@ -48,10 +48,14 @@ def generate_image(prompt,
     
     if model_choice == "14B":
         pipe = pipe_14b
-        pipe_2b.to("cpu")
+        del pipe_2b
+        torch.cuda.empty_cache()
+
+        
     else:  
         pipe = pipe_2b
-        pipe_14b.to("cpu")
+        del pipe_14b
+        torch.cuda.empty_cache()
     
     if randomize_seed:
         actual_seed = random.randint(0, 1000000)