Upload 2 files

Package versions fixes

app.py

@@ -5,28 +5,25 @@ This module provides a web UI for testing the multi-frame image restoration mode
 
 from __future__ import annotations
 
+from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
 import gradio as gr
 import numpy as np
 import torch
+import torch.nn as nn
+import torch.nn.functional as F
 from huggingface_hub import hf_hub_download
 from PIL import Image
 from torchvision import transforms
+from torchvision.ops import deform_conv2d
 
 
 # ============================================================================
 # Model Architecture (inline for single-file deployment)
 # ============================================================================
 
-from dataclasses import dataclass, field
-from typing import Literal
-
-import torch.nn as nn
-import torch.nn.functional as F
-from torchvision.ops import deform_conv2d
-
 
 @dataclass
 class FeatureFusionConfig:
@@ -34,7 +31,7 @@ class FeatureFusionConfig:
 
     in_channels: int = 3
     max_frames: int = 16
-    encoder_channels: list[int] = field(default_factory=lambda: [64, 128, 256])
+    encoder_channels: List[int] = field(default_factory=lambda: [64, 128, 256])
     encoder_blocks_per_stage: int = 2
     offset_channels: int = 64
     deform_groups: int = 8
@@ -44,14 +41,14 @@ class FeatureFusionConfig:
     decoder_blocks_per_stage: int = 2
     out_channels: int = 3
 
-    def to_dict(self) -> dict:
+    def to_dict(self) -> Dict:
         return {
             k: list(v) if isinstance(v, list) else v
             for k, v in self.__dict__.items()
         }
 
     @classmethod
-    def from_dict(cls, d: dict) -> "FeatureFusionConfig":
+    def from_dict(cls, d: Dict) -> "FeatureFusionConfig":
         return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})
 
 
@@ -79,10 +76,12 @@ class Encoder(nn.Module):
     def __init__(
         self,
         in_channels: int = 3,
-        channels: list[int] = [64, 128, 256],
+        channels: Optional[List[int]] = None,
         blocks_per_stage: int = 2,
     ):
         super().__init__()
+        if channels is None:
+            channels = [64, 128, 256]
         self.channels = channels
         self.conv_first = nn.Conv2d(in_channels, channels[0], 3, 1, 1)
         self.stages = nn.ModuleList()
@@ -213,7 +212,7 @@ class TemporalAttentionFusion(nn.Module):
         embed = embed.permute(0, 2, 1)
         return embed.unsqueeze(-1).unsqueeze(-1)
 
-    def forward(self, features: torch.Tensor, ref_idx: int | None = None) -> torch.Tensor:
+    def forward(self, features: torch.Tensor, ref_idx: Optional[int] = None) -> torch.Tensor:
         B, N, C, H, W = features.shape
         if ref_idx is None:
             ref_idx = 0
@@ -265,7 +264,7 @@ class AdaptiveFusion(nn.Module):
             ResidualBlock(channels),
         )
 
-    def forward(self, features: torch.Tensor, ref_idx: int | None = None) -> torch.Tensor:
+    def forward(self, features: torch.Tensor, ref_idx: Optional[int] = None) -> torch.Tensor:
         B, N, C, H, W = features.shape
 
         weights_list = []
@@ -292,11 +291,13 @@ class Decoder(nn.Module):
     def __init__(
         self,
         in_channels: int = 256,
-        channels: list[int] = [128, 64],
+        channels: Optional[List[int]] = None,
         out_channels: int = 3,
         blocks_per_stage: int = 2,
     ):
         super().__init__()
+        if channels is None:
+            channels = [128, 64]
 
         self.upsamples = nn.ModuleList()
         self.stages = nn.ModuleList()
@@ -372,7 +373,7 @@ class FeatureFusionModel(nn.Module):
             blocks_per_stage=config.decoder_blocks_per_stage,
         )
 
-    def forward(self, frames: torch.Tensor, ref_idx: int | None = None) -> dict[str, torch.Tensor]:
+    def forward(self, frames: torch.Tensor, ref_idx: Optional[int] = None) -> Dict[str, torch.Tensor]:
         B, N, C, H, W = frames.shape
 
         if ref_idx is None:
@@ -405,9 +406,9 @@ class FeatureFusionModel(nn.Module):
 
     def load_state_dict_with_compatibility(
         self,
-        state_dict: dict[str, torch.Tensor],
+        state_dict: Dict[str, torch.Tensor],
         strict: bool = False,
-    ) -> tuple[list[str], list[str]]:
+    ) -> Tuple[List[str], List[str]]:
         if "fusion.temporal_embed" in state_dict:
             old_embed = state_dict["fusion.temporal_embed"]
             old_num_frames = old_embed.shape[1]
@@ -444,14 +445,14 @@ class FeatureFusionModel(nn.Module):
 class ModelState:
     """Global state for loaded model."""
 
-    model: FeatureFusionModel | None = None
+    model: Optional[FeatureFusionModel] = None
     device: str = "cpu"
 
 
 state = ModelState()
 
 
-def load_model():
+def load_model() -> str:
     """Load model from Hugging Face Hub."""
     if state.model is not None:
         return "Model already loaded"
@@ -494,9 +495,9 @@ def load_model():
 
 
 def preprocess_images(
-    images: list[tuple[np.ndarray, str]] | None,
+    images: Optional[List[Tuple[np.ndarray, str]]],
     target_size: int = 256,
-) -> tuple[torch.Tensor, dict[str, Any]] | tuple[None, None]:
+) -> Tuple[Optional[torch.Tensor], Optional[Dict[str, Any]]]:
     """Preprocess input images to tensor with aspect ratio preservation."""
     if not images:
         return None, None
@@ -565,7 +566,7 @@ def tensor_to_pil(tensor: torch.Tensor) -> Image.Image:
 
 def postprocess_output(
     tensor: torch.Tensor,
-    preprocess_info: dict[str, Any],
+    preprocess_info: Dict[str, Any],
 ) -> Image.Image:
     """Remove padding from output tensor and convert to PIL Image."""
     if tensor.dim() == 4:
@@ -584,10 +585,10 @@ def postprocess_output(
 
 
 def process_images(
-    images: list[tuple[np.ndarray, str]] | None,
+    images: Optional[List[Tuple[np.ndarray, str]]],
     image_size: int = 256,
     ref_frame: int = 0,
-) -> tuple[Image.Image | None, str]:
+) -> Tuple[Optional[Image.Image], str]:
     """Process input images through the model."""
     if state.model is None:
         load_result = load_model()

requirements.txt

@@ -1,6 +1,6 @@
 torch>=2.1.0
 torchvision>=0.16.0
-gradio>=4.0.0
-huggingface_hub>=0.20.0
+gradio>=4.44.0
+huggingface_hub>=0.24.0
 numpy>=1.26.0
 Pillow>=10.0.0