description and title

app.py

@@ -9,7 +9,7 @@ import kornia.color as color
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = Generator()
-model_weights = torch.load('gen6.pth', map_location=device, weights_only=True)
+model_weights = torch.load('model.pth', map_location=device, weights_only=True)
 model.load_state_dict(model_weights)
 model = model.to(device)
 model.eval()
@@ -52,6 +52,8 @@ def predict(image):
 
 iface = gr.Interface(fn=predict,
                      inputs=gr.Image(type="pil"),
-                     outputs=gr.Image(type="pil"))
+                     outputs=gr.Image(type="pil"),
+                     title="Colorize your grayscale images",
+                     description="This model colorizes black and white images. Upload a black and white image and see the magic happen! (works best with 256x256)",)
 
 iface.launch()

model.py

@@ -8,7 +8,6 @@ class DropoutAlways(nn.Dropout2d):
     def forward(self, x):
         return F.dropout2d(x, self.p, training=True)
 
-
 class DownBlock(nn.Module):
     def __init__(self, in_channels, out_channels, normalize=True):
         super().__init__()
@@ -31,7 +30,7 @@ class UpBlock(nn.Module):
         self.block = nn.Sequential(
             nn.ConvTranspose2d(in_channels, out_channels, 4, 2, 1, bias=False if normalize else True),
             nn.InstanceNorm2d(out_channels, affine=True) if normalize else nn.Identity(),
-            DropoutAlways(p=0.5) if dropout else nn.Identity(),
+            DropoutAlways() if dropout else nn.Identity(),
             nn.ReLU() if activation == 'relu' else nn.Tanh(),
         )