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,
@@ -47,16 +47,13 @@ def generate_image(prompt,
                    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")
+        #pipe_2b.to("cpu")
     
-    else:  
-        pipe = pipe_2b
-        pipe_14b.to("cpu")
+    #else:  
+    #    pipe = pipe_2b
+    #    pipe_14b.to("cpu")
     
     if randomize_seed:
         actual_seed = random.randint(0, 1000000)