Update app.py

app.py

@@ -3,25 +3,13 @@ import spaces
 import torch
 from diffusers import Cosmos2TextToImagePipeline, EDMEulerScheduler
 import random
-
-model_14b_id = "nvidia/Cosmos-Predict2-14B-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_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 = Cosmos2TextToImagePipeline.from_pretrained(model_2b_id,
+                                                     revision="refs/pr/2",
+                                                     torch_dtype=torch.bfloat16
+                                                    )
+pipe_2b.to("cuda")
 
 scheduler = EDMEulerScheduler(
     sigma_min=0.002,
@@ -43,17 +31,9 @@ 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",
+                   model_choice="2B",
                    progress=gr.Progress(track_tqdm=True)
                   ):
-
-    if model_choice == "14B":
-        pipe = pipe_14b
-        #pipe_2b.to("cpu")
-    
-    #else:  
-    #    pipe = pipe_2b
-    #    pipe_14b.to("cpu")
     
     if randomize_seed:
         actual_seed = random.randint(0, 1000000)
@@ -62,7 +42,7 @@ def generate_image(prompt,
     
     generator = torch.Generator().manual_seed(actual_seed)
     
-    output = pipe(
+    output = pipe_2b(
         prompt=prompt,
         negative_prompt=negative_prompt,
         generator=generator
@@ -86,7 +66,7 @@ example_prompts = [
 with gr.Blocks() as demo:
     gr.Markdown(
         """
-        # Cosmos-Predict2 14B Text2Image
+        # Cosmos-Predict2 2B Text2Image
         [[Model]](https://huggingface.co/nvidia/Cosmos-Predict2-14B-Text2Image), [[Code]](https://github.com/nvidia-cosmos/cosmos-predict2)
         """
     )
@@ -121,7 +101,7 @@ with gr.Blocks() as demo:
                 
             model_radio = gr.Radio(
                 choices=["14B", "2B"],
-                value="14B",
+                value="2B",
                 label="Model Selection",
                 visible=False
             )