Update app.py

app.py

@@ -12,16 +12,16 @@ pipe_14b = Cosmos2TextToImagePipeline.from_pretrained(model_14b_id,
                                                       torch_dtype=torch.bfloat16
                                                      )
 
-pipe_2b = Cosmos2TextToImagePipeline.from_pretrained(model_2b_id,
-                                                     revision="refs/pr/2",
-                                                     text_encoder=pipe_14b.text_encoder,
-                                                     safety_checker=pipe_14b.safety_checker,
-                                                     vae=pipe_14b.vae,
-                                                     tokenizer=pipe_14b.tokenizer,
-                                                     torch_dtype=torch.bfloat16
-                                                    )
+#pipe_2b = Cosmos2TextToImagePipeline.from_pretrained(model_2b_id,
+#                                                     revision="refs/pr/2",
+#                                                     text_encoder=pipe_14b.text_encoder,
+#                                                     safety_checker=pipe_14b.safety_checker,
+#                                                     vae=pipe_14b.vae,
+#                                                     tokenizer=pipe_14b.tokenizer,
+#                                                     torch_dtype=torch.bfloat16
+#                                                    )
 pipe_14b.to("cuda")
-pipe_2b.to("cuda")
+#pipe_2b.to("cuda")
 
 scheduler = EDMEulerScheduler(
     sigma_min=0.002,
@@ -46,16 +46,17 @@ def generate_image(prompt,
                    model_choice="14B",
                    progress=gr.Progress(track_tqdm=True)
                   ):
+
+    pipe_14b.to("cuda")
+    pipe_2b.to("cuda")
     
     if model_choice == "14B":
         pipe = pipe_14b
         pipe_2b.to("cpu")
-        torch.cuda.empty_cache()
-
+    
     else:  
         pipe = pipe_2b
         pipe_14b.to("cpu")
-        torch.cuda.empty_cache()
     
     if randomize_seed:
         actual_seed = random.randint(0, 1000000)