Update app.py

app.py

@@ -109,7 +109,7 @@ def load_and_prepare_model(model_id):
         "ford442/RealVisXL_V5.0_BF16": torch.bfloat16,
     }
     dtype = model_dtypes.get(model_id, torch.bfloat16)  # Default to float32 if not found
-    vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16", torch_dtype=torch.bfloat16,safety_checker=None).to(torch.bfloat16).to('cuda')
+    vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16", torch_dtype=torch.bfloat16,safety_checker=None)
     pipe = StableDiffusionXLPipeline.from_pretrained(
         model_id,
         torch_dtype=torch.bfloat16,
@@ -149,13 +149,12 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-    
-
 GPU_DURATION_OPTIONS = {
+    "Short (25s)": 25,
     "Short (45s)": 45,
-    "Short (60s)": 60,
+    "Medium (60s)": 60,
     "Medium (80s)": 80,
-    "Medium (100s)": 100,
+    "Long (100s)": 100,
     "Long (120s)": 120,
     "Long (140s)": 140,
 }
@@ -163,7 +162,7 @@ GPU_DURATION_OPTIONS = {
 def set_gpu_duration(duration_choice):
     os.environ["GPU_DURATION"] = str(GPU_DURATION_OPTIONS[duration_choice])
 
-@spaces.GPU(duration=int(os.getenv("GPU_DURATION", "80")))  # Dynamic duration
+@spaces.GPU(duration=int(os.getenv("GPU_DURATION", "60")))  # Dynamic duration
 def generate(
     model_choice: str,
     prompt: str,
@@ -320,7 +319,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         gpu_duration = gr.Dropdown(
             label="GPU Duration",
             choices=list(GPU_DURATION_OPTIONS.keys()),
-            value="Medium (80s)"  # Default value
+            value="Medium (60s)"  # Default value
         )
         
         gpu_duration.change(fn=set_gpu_duration, inputs=gpu_duration, outputs=[])