Update app.py

app.py

@@ -6,44 +6,23 @@ import random
 
 model_14b_id = "nvidia/Cosmos-Predict2-14B-Text2Image"
 
-pipe_14b = Cosmos2TextToImagePipeline.from_pretrained(model_14b_id,
-                                                      torch_dtype=torch.bfloat16
-                                                     )
-
-pipe_14b.to("cuda")
-
-scheduler = EDMEulerScheduler(
-    sigma_min=0.002,
-    sigma_max=80,
-    sigma_data=1.0,
-    sigma_schedule="karras",
-    num_train_timesteps=1000,
-    prediction_type="epsilon",
-    rho=7.0,
-    final_sigmas_type="sigma_min",
-    # use_flow_sigmas=True,
+pipe_14b = Cosmos2TextToImagePipeline.from_pretrained(
+    model_14b_id,
+    torch_dtype=torch.bfloat16
 )
 
-pipe_14b.scheduler = scheduler
-# pipe_2b.scheduler = scheduler
+pipe_14b.to("cuda")
 
 @spaces.GPU(duration=120)
-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)
-                  ):
+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_2b.to("cpu")
-    
-    #else:  
-    #    pipe = pipe_2b
-    #    pipe_14b.to("cpu")
-    
     if randomize_seed:
         actual_seed = random.randint(0, 1000000)
     else:
@@ -51,7 +30,7 @@ def generate_image(prompt,
     
     generator = torch.Generator().manual_seed(actual_seed)
     
-    output = pipe(
+    output = pipe_14b(
         prompt=prompt,
         negative_prompt=negative_prompt,
         generator=generator
@@ -59,8 +38,6 @@ def generate_image(prompt,
     
     return output, actual_seed
 
-
-
 example_prompts = [
     "A well-worn broom sweeps across a dusty wooden floor, its bristles gathering crumbs and flecks of debris in swift, rhythmic strokes. Dust motes dance in the sunbeams filtering through the window, glowing momentarily before settling. The quiet swish of straw brushing wood is interrupted only by the occasional creak of old floorboards. With each pass, the floor grows cleaner, restoring a sense of quiet order to the humble room.",
     "A laundry machine whirs to life, tumbling colorful clothes behind the foggy glass door. Suds begin to form in a frothy dance, clinging to fabric as the drum spins. The gentle thud of shifting clothes creates a steady rhythm, like a heartbeat of the home. Outside the machine, a quiet calm fills the room, anticipation building for the softness and warmth of freshly laundered garments.",