Update app.py

app.py

@@ -5,21 +5,21 @@ from diffusers import Cosmos2TextToImagePipeline, EDMEulerScheduler
 import random
 
 model_14b_id = "nvidia/Cosmos-Predict2-14B-Text2Image"
-# model_2b_id = "nvidia/Cosmos-Predict2-2B-Text2Image"
+model_2b_id = "nvidia/Cosmos-Predict2-2B-Text2Image"
 
 pipe_14b = Cosmos2TextToImagePipeline.from_pretrained(model_14b_id,
                                                       revision="refs/pr/1",
                                                       torch_dtype=torch.bfloat16
                                                      )
 pipe_14b.to("cuda")
-#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
+                                                    )
 
 scheduler = EDMEulerScheduler(
     sigma_min=0.002,
@@ -39,11 +39,11 @@ pipe_14b.scheduler = scheduler
 @spaces.GPU
 def generate_image(prompt, negative_prompt="The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality.", seed=42, randomize_seed=False, model_choice="14B", progress=gr.Progress(track_tqdm=True)):
     
-    #if model_choice == "14B":
-    #    pipe = pipe_14b
-    #else:  
-    #    pipe = pipe_2b
-    #    pipe.to("cuda")
+    if model_choice == "14B":
+        pipe = pipe_14b
+    else:  
+        pipe = pipe_2b
+        pipe.to("cuda")
     
     if randomize_seed:
         actual_seed = random.randint(0, 1000000)
@@ -52,7 +52,7 @@ def generate_image(prompt, negative_prompt="The video captures a series of frame
     
     generator = torch.Generator().manual_seed(actual_seed)
     
-    output = pipe_14b(
+    output = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
         generator=generator