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_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 = 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,
@@ -34,17 +34,16 @@ scheduler = EDMEulerScheduler(
 )
 
 pipe_14b.scheduler = scheduler
-pipe_2b.scheduler = scheduler
+# 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)):
     
-    if model_choice == "14B":
-        pipe = pipe_14b
-        pipe.to("cuda")
-    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,9 +51,13 @@ def generate_image(prompt, negative_prompt="The video captures a series of frame
         actual_seed = seed
     
     generator = torch.Generator().manual_seed(actual_seed)
-    output = pipe(
-        prompt=prompt, negative_prompt=negative_prompt, generator=generator
+    
+    output = pipe_14b(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        generator=generator
     ).images[0]
+    
     return output, actual_seed
 
 
@@ -73,10 +76,7 @@ example_prompts = [
 with gr.Blocks() as demo:
     gr.Markdown(
         """
-        # Cosmos Predict-2 Text-to-Image Generator
-        Enter a detailed prompt to generate an image using the Cosmos model.
-        You can also provide a negative prompt to guide the generation away from certain elements.
-        Choose between the 14B model (higher quality, slower) or 2B model (faster, smaller).
+        # Cosmos Predict-2 14B Text-to-Image Generator
         """
     )
     with gr.Row():
@@ -112,6 +112,7 @@ with gr.Blocks() as demo:
                 choices=["14B", "2B"],
                 value="14B",
                 label="Model Selection",
+                visible=False
             )
             generate_button = gr.Button("Generate Image")