Do not assign gpu directly

models/utils.py

@@ -36,7 +36,7 @@ def get_model(
             cache_dir=cache_dir,
             memsave=memsave,
         )
-        pipe = pipe.to(device, dtype)
+        #pipe = pipe.to(device, dtype)
     elif model_name == "sdxl-turbo":
         vae = AutoencoderKL.from_pretrained(
             "madebyollin/sdxl-vae-fp16-fix",
@@ -55,7 +55,7 @@ def get_model(
         pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(
             pipe.scheduler.config, timestep_spacing="trailing"
         )
-        pipe = pipe.to(device, dtype)
+        #pipe = pipe.to(device, dtype)
     elif model_name == "pixart":
         pipe = RewardPixartPipeline.from_pretrained(
             "PixArt-alpha/PixArt-XL-2-1024-MS",
@@ -80,7 +80,7 @@ def get_model(
         pipe.transformer.eval()
         freeze_params(pipe.transformer.parameters())
         pipe.transformer.enable_gradient_checkpointing()
-        pipe = pipe.to(device)
+        #pipe = pipe.to(device)
     elif model_name == "hyper-sd":
         base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
         repo_name = "ByteDance/Hyper-SD"
@@ -108,20 +108,20 @@ def get_model(
         pipe.scheduler = LCMScheduler.from_config(
             pipe.scheduler.config, cache_dir=cache_dir
         )
-        pipe = pipe.to(device, dtype)
+        #pipe = pipe.to(device, dtype)
         # upcast vae
         pipe.vae = pipe.vae.to(dtype=torch.float32)
     elif model_name == "flux":
         pipe = RewardFluxPipeline.from_pretrained(
             "black-forest-labs/FLUX.1-schnell",
-            torch_dtype=torch.bfloat16,
+            torch_dtype=torch.float16,
             cache_dir=cache_dir,
         )
-        pipe.to(device, dtype)
+        #pipe.to(device, dtype)
     else:
         raise ValueError(f"Unknown model name: {model_name}")
-    if enable_sequential_cpu_offload:
-        pipe.enable_sequential_cpu_offload()
+    #if enable_sequential_cpu_offload:
+    #    pipe.enable_sequential_cpu_offload()
     return pipe