Update app.py

app.py

@@ -30,13 +30,8 @@ USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
 ENABLE_REFINER = os.getenv("ENABLE_REFINER", "0") == "1"
 
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
 
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 if torch.cuda.is_available():
     vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
     pipe = DiffusionPipeline.from_pretrained("dataautogpt3/OpenDalleV1.1", vae=vae, torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
@@ -56,6 +51,11 @@ if torch.cuda.is_available():
         if ENABLE_REFINER:
             refiner.unet = torch.compile(refiner.unet, mode="reduce_overhead", fullgraph=True)
 
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
 @spaces.GPU(enable_queue=True)
 def infer(
     prompt: str,