Make the generator CPU

app.py

@@ -160,7 +160,7 @@ def save_slim_model(infinity_model_path, save_file=None, device='cpu', key='gpt_
     return save_file
 
 def load_tokenizer(t5_path =''):
-    print(f'[Loading tokenizer and text encoder]')
+    print('[Loading tokenizer and text encoder]')
     text_tokenizer: T5TokenizerFast = AutoTokenizer.from_pretrained(t5_path, revision=None, legacy=True)
     text_tokenizer.model_max_length = 512
     text_encoder: T5EncoderModel = T5EncoderModel.from_pretrained(t5_path, torch_dtype=torch.float16)
@@ -186,7 +186,7 @@ def load_infinity(
     use_flex_attn=False,
     bf16=False,
 ):
-    print(f'[Loading Infinity]')
+    print('[Loading Infinity]')
     
     # Set device if not provided
     if device is None:
@@ -232,13 +232,13 @@ def load_infinity(
         infinity_test.eval()
         infinity_test.requires_grad_(False)
 
-        print(f'[Load Infinity weights]')
+        print('[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)
-        
+        infinity_test.rng = torch.Generator(device="cpu")
+
         return infinity_test
 
 def transform(pil_img, tgt_h, tgt_w):