5tasks

app.py

@@ -12,8 +12,50 @@ from net.mask_spatial_mixer_channel_easyshallow_harddeep_twomore_prompt import S
 from einops import rearrange
 import spaces
 
-from huggingface_hub import PyTorchModelHubMixin
 
+class SpatialMixerChannelEasyShallowHardDeepTwoMorePromptIRModel(pl.LightningModule):
+    def __init__(self):
+        super().__init__()
+        
+        self.net =  SpatialMixerChannelEasyShallowHardDeepTwoMorePrompt(
+                    inp_channels=3,
+                    out_channels=3,
+                    dim = 48,
+                    num_blocks = [2,4,4,4],
+                    num_refinement_blocks = 4,
+                    channel_heads = [1,1,1,1],
+                    spatial_heads = [1,2,4,8],
+                    overlap_ratio = 0.5,
+                    dim_head = 16,
+                    ratio = 0.5,
+                    window_size = 8,
+                    bias = False,
+                    ffn_expansion_factor = 2.66,
+                    LayerNorm_type = 'WithBias',   ## Other option 'BiasFree'
+                    dual_pixel_task = False,        ## True for dual-pixel defocus deblurring only. Also set inp_channels=6
+                    scale = 1,
+                    prompt = True,
+                    hard_ratio = 0.5
+                    )
+        self.loss_fn  = nn.L1Loss()
+        
+        # Initialize EMA
+        self.ema_decay = 0.999
+        self.ema_model = self.create_ema_model()
+        
+    def create_ema_model(self):
+        ema_model = copy.deepcopy(self.net)
+        for param in ema_model.parameters():
+            param.detach_()
+        return ema_model
+    
+    def update_ema_model(self):
+        with torch.no_grad():
+            for ema_param, model_param in zip(self.ema_model.parameters(), self.net.parameters()):
+                ema_param.data = self.ema_decay * ema_param.data + (1 - self.ema_decay) * model_param.data
+    
+    def forward(self,x, training = False):
+        return self.net(x, training = training)
 
 
 # crop an image to the multiple of base
@@ -23,20 +65,6 @@ def crop_img(image, base=64):
     crop_h = h % base
     crop_w = w % base
     return image[crop_h // 2:h - crop_h + crop_h // 2, crop_w // 2:w - crop_w + crop_w // 2, :]
-
-class CatAIR5tasksModel(
-    nn.Module,
-    PyTorchModelHubMixin,
-    repo_url="https://huggingface.co/jiachenjiang/Cat-AIR-5tasks",
-    pipeline_tag="image-to-image",
-    license="mit",
-):
-    def __init__(self, **kwargs):
-        super().__init__()
-        self.net = SpatialMixerChannelEasyShallowHardDeepTwoMorePrompt(**kwargs)
-
-    def forward(self, x):
-        return self.net(x)
     
 def np_to_pil(img_np):
     """
@@ -73,8 +101,14 @@ def restore_image(input_img):
     np.random.seed(0)
     torch.manual_seed(0)
 
-    model = CatAIR5tasksModel.from_pretrained("jiachenjiang/Cat-AIR-5tasks").cuda()
-    net = model.net.eval()
+
+    ckpt_path = "/local/storage1/jiang.2880/Cat-AIR/ckpt/epoch=71-step=228600.ckpt"
+    print("CKPT name : {}".format(ckpt_path))
+
+    checkpoint = torch.load(ckpt_path, map_location='cuda', weights_only=True)
+    model = SpatialMixerChannelEasyShallowHardDeepTwoMorePromptIRModel()
+    model.net.load_state_dict(checkpoint['ema_model_state_dict'])
+    net = model.net.cuda()
 
 
     degraded_img = crop_img(np.array(input_img.convert('RGB')), base=16)

ckpt/epoch=71-step=228600.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1586ac517148dc24dd1588af6aa11205f8628253c4b7719211bd7af6568fb325
+size 535261296

requirements.txt

@@ -7,5 +7,4 @@ pytorch-lightning
 pydantic
 pillow
 einops
-numpy<2
-huggingface_hub
+numpy<2