Update app.py

app.py

@@ -11,16 +11,16 @@ 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_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,
@@ -37,13 +37,21 @@ scheduler = EDMEulerScheduler(
 pipe_14b.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)):
+@spaces.GPU(duration=75)
+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_2b.to("cpu")
+    else:  
+        pipe = pipe_2b
+        pipe_14b.to("cpu")
     
     if randomize_seed:
         actual_seed = random.randint(0, 1000000)