Make device optional in load_infinity function; set default to 'cuda' or 'cpu' based on availability and adjust autocast dtype handling

app.py

@@ -188,7 +188,7 @@ def load_infinity(
     model_path='', 
     scale_schedule=None, 
     vae=None, 
-    device='cuda',
+    device=None,  # Make device optional
     model_kwargs=None,
     text_channels=2048,
     apply_spatial_patchify=0,
@@ -196,9 +196,23 @@ def load_infinity(
     bf16=False,
 ):
     print(f'[Loading Infinity]')
+    
+    # Set device if not provided
+    if device is None:
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    print(f'Using device: {device}')
+
+    # Set autocast dtype based on bf16 and device support
+    if bf16 and device == 'cuda' and torch.cuda.is_bf16_supported():
+        autocast_dtype = torch.bfloat16
+    else:
+        autocast_dtype = torch.float32
+        bf16 = False  # Disable bf16 if not supported
+
     text_maxlen = 512
     torch.cuda.empty_cache()
-    with torch.amp.autocast(device_type=device, enabled=True, dtype=torch.bfloat16, cache_enabled=True), torch.no_grad():
+
+    with torch.amp.autocast(device_type=device, enabled=bf16, dtype=autocast_dtype, cache_enabled=True), torch.no_grad():
         infinity_test: Infinity = Infinity(
             vae_local=vae, text_channels=text_channels, text_maxlen=text_maxlen,
             shared_aln=True, raw_scale_schedule=scale_schedule,
@@ -217,6 +231,7 @@ def load_infinity(
             train_h_div_w_list=[1.0],
             **model_kwargs,
         ).to(device)
+
         print(f'[you selected Infinity with {model_kwargs=}] model size: {sum(p.numel() for p in infinity_test.parameters())/1e9:.2f}B, bf16={bf16}')
 
         if bf16:
@@ -229,7 +244,10 @@ def load_infinity(
         print(f'[Load Infinity weights]')
         state_dict = torch.load(model_path, map_location=device)
         print(infinity_test.load_state_dict(state_dict))
+        
+        # Initialize random number generator on the correct device
         infinity_test.rng = torch.Generator(device=device)
+        
         return infinity_test
 
 def transform(pil_img, tgt_h, tgt_w):