Update src/models/sd2_sr.py

src/models/sd2_sr.py

@@ -187,6 +187,7 @@ def load_model(dtype=torch.bfloat16, device=None):
     encoder.to(dtype=dtype, device=device)
     encoder.device = device
 
+    # Create a new DDIM instance with the updated components
     ddim = DDIM(config, vae, encoder, unet)
 
     params = {
@@ -205,4 +206,13 @@ def load_model(dtype=torch.bfloat16, device=None):
     low_scale_model = low_scale_model.to(dtype=dtype, device=device)
 
     ddim.low_scale_model = low_scale_model
-    return ddim
+    
+    # Add a 'to' method to the DDIM object to make it compatible with the patching code
+    def ddim_to(self, *args, **kwargs):
+        # This is a no-op function since the components are already moved to the correct device
+        return self
+    
+    # Add the 'to' method to the DDIM instance
+    ddim.to = ddim_to.__get__(ddim)
+    
+    return ddim