Ensure model components are on CPU after loading in app.py to prevent device-related issues during processing.

app.py

@@ -146,6 +146,13 @@ model = CRM(specs)
 model.load_state_dict(torch.load(crm_path, map_location="cpu"), strict=False)
 model = model.to(args.device)
 
+# After loading or instantiating the model, ensure everything is on CPU
+model = model.cpu()
+if hasattr(model, 'rgbMlp'):
+    model.rgbMlp = model.rgbMlp.cpu()
+if hasattr(model, 'decoder'):
+    model.decoder = model.decoder.cpu()
+
 stage1_config = OmegaConf.load(args.stage1_config).config
 stage2_config = OmegaConf.load(args.stage2_config).config
 stage2_sampler_config = stage2_config.sampler