Do not load models on gpu at first

models/utils.py

@@ -81,36 +81,41 @@ def get_model(
         freeze_params(pipe.transformer.parameters())
         pipe.transformer.enable_gradient_checkpointing()
         #pipe = pipe.to(device)
+    
     elif model_name == "hyper-sd":
         base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
         repo_name = "ByteDance/Hyper-SD"
         ckpt_name = "Hyper-SDXL-1step-Unet.safetensors"
-        # Load model.
+    
+        # Load model but don't specify device or dtype (defaults to CPU and float32)
         unet = UNet2DConditionModel.from_config(
             base_model_id, subfolder="unet", cache_dir=cache_dir
-        ).to(device, dtype)
+        )
+    
+        # Load state dict into unet (stays on CPU by default)
         unet.load_state_dict(
             load_file(
                 hf_hub_download(repo_name, ckpt_name, cache_dir=cache_dir),
                 device="cuda",
             )
         )
+    
+        # Initialize the pipeline (it will stay on CPU initially, using default dtype)
         pipe = RewardStableDiffusionXL.from_pretrained(
             base_model_id,
             unet=unet,
-            torch_dtype=dtype,
-            variant="fp16",
+            torch_dtype=torch.float16,
+            variant="fp16",  # Still set fp16 for later use on GPU
             cache_dir=cache_dir,
             is_hyper=True,
             memsave=memsave,
         )
+
         # Use LCM scheduler instead of ddim scheduler to support specific timestep number inputs
         pipe.scheduler = LCMScheduler.from_config(
             pipe.scheduler.config, cache_dir=cache_dir
         )
-        #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",
@@ -187,4 +192,4 @@ def get_multi_apply_fn(
             generator=generator,
         )
     else:
-        raise ValueError(f"Unknown model type: {model_type}")
+        raise ValueError(f"Unknown model type: {model_type}")