Update app.py

app.py

@@ -7,11 +7,19 @@ import random
 model_14b_id = "diffusers-internal-dev/ct2i-14B"
 model_2b_id = "diffusers-internal-dev/ct2i-2B"
 
-pipe_14b = Cosmos2TextToImagePipeline.from_pretrained(model_14b_id, torch_dtype=torch.bfloat16)
-pipe_14b.to("cuda")
+pipe_14b = Cosmos2TextToImagePipeline.from_pretrained(model_14b_id,
+                                                      revision="refs/pr/1",
+                                                      torch_dtype=torch.bfloat16
+                                                     )
 
-pipe_2b = Cosmos2TextToImagePipeline.from_pretrained(model_2b_id, 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.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
+                                                    )
 
 scheduler = EDMEulerScheduler(
     sigma_min=0.002,
@@ -30,11 +38,13 @@ pipe_2b.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)):
-    # Select the appropriate pipeline based on model choice
+    
     if model_choice == "14B":
         pipe = pipe_14b
-    else:  # "2B Model"
+        pipe.to("cuda")
+    else:  
         pipe = pipe_2b
+        pipe.to("cuda")
     
     if randomize_seed:
         actual_seed = random.randint(0, 1000000)