Update app.py

app.py

@@ -13,12 +13,7 @@ pipe = SanaSprintPipeline.from_pretrained(
     "Efficient-Large-Model/Sana_Sprint_0.6B_1024px_diffusers",
     torch_dtype=torch.bfloat16
 )
-pipe2 = SanaSprintPipeline.from_pretrained(
-    "Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers",
-    torch_dtype=torch.bfloat16
-)
 pipe.to(device)
-pipe2.to(device)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
@@ -29,7 +24,7 @@ def infer(prompt, model_size, seed=42, randomize_seed=False, width=1024, height=
     generator = torch.Generator().manual_seed(seed)
     
     # Choose the appropriate model based on selected model size
-    selected_pipe = pipe if model_size == "0.6B" else pipe2
+    selected_pipe = pipe
     
     img = selected_pipe(
             prompt=prompt,
@@ -44,8 +39,8 @@ def infer(prompt, model_size, seed=42, randomize_seed=False, width=1024, height=
     return img.images[0], seed
     
 examples = [
-    ["a tiny astronaut hatching from an egg on the moon", "1.6B"],
-    ["🐶 Wearing 🕶 flying on the 🌈", "1.6B"],
+    ["a tiny astronaut hatching from an egg on the moon", "0.6B"],
+    ["🐶 Wearing 🕶 flying on the 🌈", "0.6B"],
     ["an anime illustration of a wiener schnitzel", "0.6B"],
     ["a photorealistic landscape of mountains at sunset", "0.6B"],
 ]
@@ -78,8 +73,8 @@ with gr.Blocks(css=css) as demo:
         
         model_size = gr.Radio(
             label="Model Size",
-            choices=["0.6B", "1.6B"],
-            value="1.6B",
+            choices=["0.6B"],
+            value="0.6B",
             interactive=True
         )