Update app.py

app.py

@@ -124,7 +124,7 @@ def load_and_prepare_model(model_id):
 # Preload and compile both models
 models = {key: load_and_prepare_model(value) for key, value in MODEL_OPTIONS.items()}
 
-MAX_SEED = np.iinfo(np.int64).max
+MAX_SEED = np.iinfo(np.int32).max
 
 def upload_to_ftp(filename):
     try:
@@ -148,8 +148,22 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     return seed
+
     
-@spaces.GPU(duration=60)
+
+GPU_DURATION_OPTIONS = {
+    "Short (45s)": 45,
+    "Short (60s)": 60,
+    "Medium (80s)": 80,
+    "Medium (100s)": 100,
+    "Long (120s)": 120,
+    "Long (140s)": 140,
+}
+
+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
 def generate(
     model_choice: str,
     prompt: str,
@@ -169,7 +183,7 @@ def generate(
     global models
     pipe = models[model_choice]
     seed = int(randomize_seed_fn(seed, randomize_seed))
-    generator = torch.Generator(device='cuda').manual_seed(seed)
+    generator = torch.Generator(device='cpu').manual_seed(seed)
 
     prompt, negative_prompt = apply_style(style_selection, prompt, negative_prompt)
 
@@ -188,7 +202,8 @@ def generate(
         options["use_resolution_binning"] = True
 
     images = []
-    for i in range(0, num_images, BATCH_SIZE):
+    with torch.no_grad():
+      for i in range(0, num_images, BATCH_SIZE):
         batch_options = options.copy()
         batch_options["prompt"] = options["prompt"][i:i+BATCH_SIZE]
         if "negative_prompt" in batch_options: