move torch.Generator from switti init to pipeline call

app.py

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

models/pipeline.py

@@ -113,8 +113,7 @@ class SwittiPipeline:
         if seed is None:
             rng = None
         else:
-            switti.rng.manual_seed(seed)
-            rng = switti.rng
+            rng = torch.Generator(self.device).manual_seed(seed)
 
         context, cond_vector, context_attn_bias = self.encode_prompt(prompt, null_prompt)
 

models/switti.py

@@ -72,7 +72,6 @@ class Switti(nn.Module):
         self.rope = rope
 
         self.num_stages_minus_1 = len(self.patch_nums) - 1
-        self.rng = torch.Generator(device=device)
 
         # 1. input (word) embedding
         self.word_embed = nn.Linear(self.Cvae, self.C)