Update app.py

app.py

@@ -11,6 +11,7 @@ pipe_14b = Cosmos2TextToImagePipeline.from_pretrained(model_14b_id,
                                                       revision="refs/pr/1",
                                                       torch_dtype=torch.bfloat16
                                                      )
+
 pipe_2b = Cosmos2TextToImagePipeline.from_pretrained(model_2b_id,
                                                      revision="refs/pr/2",
                                                      text_encoder=pipe_14b.text_encoder,
@@ -35,7 +36,7 @@ scheduler = EDMEulerScheduler(
 )
 
 pipe_14b.scheduler = scheduler
-#pipe_2b.scheduler = scheduler
+pipe_2b.scheduler = scheduler
 
 @spaces.GPU(duration=75)
 def generate_image(prompt,
@@ -48,13 +49,12 @@ def generate_image(prompt,
     
     if model_choice == "14B":
         pipe = pipe_14b
-        del pipe_2b
+        pipe_2b.to("cpu")
         torch.cuda.empty_cache()
 
-        
     else:  
         pipe = pipe_2b
-        del pipe_14b
+        pipe_14b.to("cpu")
         torch.cuda.empty_cache()
     
     if randomize_seed: