Update inference_coz.py

inference_coz.py

@@ -175,11 +175,11 @@ if __name__ == "__main__":
         if not args.efficient_memory:
             from osediff_sd3 import OSEDiff_SD3_TEST, SD3Euler
             model = SD3Euler()
-            model.text_enc_1.to('cuda:0')
-            model.text_enc_2.to('cuda:0')
-            model.text_enc_3.to('cuda:0')
-            model.transformer.to('cuda:1', dtype=torch.float32)
-            model.vae.to('cuda:1', dtype=torch.float32)
+            model.text_enc_1.to('cuda')
+            model.text_enc_2.to('cuda')
+            model.text_enc_3.to('cuda')
+            model.transformer.to('cuda', dtype=torch.float32)
+            model.vae.to('cuda', dtype=torch.float32)
             for p in [model.text_enc_1, model.text_enc_2, model.text_enc_3, model.transformer, model.vae]:
                 p.requires_grad_(False)
             model_test = OSEDiff_SD3_TEST(args, model)
@@ -335,9 +335,9 @@ if __name__ == "__main__":
                 if args.efficient_memory and model is not None:
                     print("Ensuring SR model components are on CUDA for SR inference.")
                     if not isinstance(model_test, OSEDiff_SD3_TEST_efficient):
-                        model.text_enc_1.to('cuda:0')
-                        model.text_enc_2.to('cuda:0')
-                        model.text_enc_3.to('cuda:0')
+                        model.text_enc_1.to('cuda')
+                        model.text_enc_2.to('cuda')
+                        model.text_enc_3.to('cuda')
                     # transformer and VAE should already be on CUDA per initialization
                     model.transformer.to('cuda', dtype=torch.float32)
                     model.vae.to('cuda', dtype=torch.float32)