update

asset3d_gen/models/sr_model.py

@@ -59,55 +59,58 @@ class ImageStableSR:
 
 class ImageRealESRGAN:
     def __init__(self, outscale: int, model_path: str = None) -> None:
-        # monkey_patch
+        # monkey patch to support torchvision>=0.16
         import torchvision
         from packaging import version
 
         if version.parse(torchvision.__version__) > version.parse("0.16"):
             import sys
             import types
-
             import torchvision.transforms.functional as TF
 
-            functional_tensor = types.ModuleType(
-                "torchvision.transforms.functional_tensor"
-            )
+            functional_tensor = types.ModuleType("torchvision.transforms.functional_tensor")
             functional_tensor.rgb_to_grayscale = TF.rgb_to_grayscale
-            sys.modules["torchvision.transforms.functional_tensor"] = (
-                functional_tensor
-            )
-
-        from basicsr.archs.rrdbnet_arch import RRDBNet
-        from realesrgan import RealESRGANer
+            sys.modules["torchvision.transforms.functional_tensor"] = functional_tensor
 
         self.outscale = outscale
-        model = RRDBNet(
-            num_in_ch=3,
-            num_out_ch=3,
-            num_feat=64,
-            num_block=23,
-            num_grow_ch=32,
-            scale=4,
-        )
-        if model_path is None:
-            suffix = "super_resolution"
-            model_path = snapshot_download(
-                repo_id="xinjjj/RoboAssetGen", allow_patterns=f"{suffix}/*"
-            )
-            model_path = os.path.join(
-                model_path, suffix, "RealESRGAN_x4plus.pth"
+        self.model_path = model_path
+        self.upsampler = None
+        
+    def _lazy_init(self):
+        if self.upsampler is None:
+            from basicsr.archs.rrdbnet_arch import RRDBNet
+            from realesrgan import RealESRGANer
+            from huggingface_hub import snapshot_download
+
+            model = RRDBNet(
+                num_in_ch=3,
+                num_out_ch=3,
+                num_feat=64,
+                num_block=23,
+                num_grow_ch=32,
+                scale=4,
             )
 
-        self.upsampler = RealESRGANer(
-            scale=4,
-            model_path=model_path,
-            model=model,
-            pre_pad=0,
-            half=True,
-        )
+            model_path = self.model_path
+            if model_path is None:
+                suffix = "super_resolution"
+                model_path = snapshot_download(
+                    repo_id="xinjjj/RoboAssetGen", allow_patterns=f"{suffix}/*"
+                )
+                model_path = os.path.join(model_path, suffix, "RealESRGAN_x4plus.pth")
+
+            self.upsampler = RealESRGANer(
+                scale=4,
+                model_path=model_path,
+                model=model,
+                pre_pad=0,
+                half=True,
+            )
 
     @spaces.GPU
     def __call__(self, image: Union[Image.Image, np.ndarray]) -> Image.Image:
+        self._lazy_init()
+        
         if isinstance(image, Image.Image):
             image = np.array(image)