remove torch.cuda.is_available from inner scripts

app.py

@@ -7,6 +7,7 @@ from models import SwittiPipeline
 import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
+print("Device:", device)
 model_repo_id = "yresearch/Switti"
 
 

models/switti.py

@@ -53,6 +53,7 @@ class Switti(nn.Module):
         use_swiglu_ffn=True,
         use_ar=False,
         use_crop_cond=True,
+        device='cuda',
     ):
         super().__init__()
         # 0. hyperparameters
@@ -71,7 +72,7 @@ class Switti(nn.Module):
         self.rope = rope
 
         self.num_stages_minus_1 = len(self.patch_nums) - 1
-        self.rng = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
+        self.rng = torch.Generator(device=device)
 
         # 1. input (word) embedding
         self.word_embed = nn.Linear(self.Cvae, self.C)
@@ -392,6 +393,7 @@ class SwittiHF(Switti, PyTorchModelHubMixin):
         use_swiglu_ffn=True,
         use_ar=False,
         use_crop_cond=True,
+        device='cuda',
     ):
         heads = depth
         width = depth * 64
@@ -406,4 +408,5 @@ class SwittiHF(Switti, PyTorchModelHubMixin):
             use_swiglu_ffn=use_swiglu_ffn,
             use_ar=use_ar,
             use_crop_cond=use_crop_cond,
+            device=device,
         )