imports

app.py

@@ -7,7 +7,7 @@ import gradio as gr
 import albumentations as A
 from albumentations.pytorch import ToTensorV2
 from PIL import Image
-import transformers
+from transformers import ConvNextV2Config, ConvNextV2ForImageClassification
 
 # preprocessing
 transforms = A.Compose([
@@ -22,12 +22,11 @@ transforms = A.Compose([
 class PrHu_model(nn.Module):
     def __init__(self):
         super().__init__()
-        self.configuration = transformers.ConvNextV2Config(num_channels=1, drop_path_rate=0, image_size=384, num_labels=1,
+        self.configuration = ConvNextV2Config(num_channels=1, drop_path_rate=0, image_size=384, num_labels=1,
                                                            depths=[2, 2, 6, 2], hidden_sizes=[16, 32, 64, 128])
-        self.model = transformers.ConvNextV2ForImageClassification(self.configuration)
+        self.model = ConvNextV2ForImageClassification(self.configuration)
 
     def forward(self, x):
-        # print ('starting model F pass')
         return self.model(x).logits
     
 device = 'cuda' if torch.cuda.is_available() else 'cpu'