Ensure additional model components are moved to CPU in app.py to maintain consistency and prevent device-related issues during processing.

app.py

@@ -152,6 +152,12 @@ if hasattr(model, 'rgbMlp'):
     model.rgbMlp = model.rgbMlp.cpu()
 if hasattr(model, 'decoder'):
     model.decoder = model.decoder.cpu()
+if hasattr(model, 'unet2'):
+    model.unet2 = model.unet2.cpu()
+    if hasattr(model.unet2, 'unet'):
+        model.unet2.unet = model.unet2.unet.cpu()
+if hasattr(model, 'lora'):
+    model.lora = model.lora.cpu()
 
 stage1_config = OmegaConf.load(args.stage1_config).config
 stage2_config = OmegaConf.load(args.stage2_config).config