Update app.py

app.py

@@ -8,7 +8,6 @@ import numpy as np
 import requests
 from io import BytesIO
 
-# Download the model from Hugging Face Hub
 repo_id = "Hammad712/GAN-Colorization-Model"
 model_filename = "generator.pt"
 model_path = hf_hub_download(repo_id=repo_id, filename=model_filename)
@@ -29,7 +28,6 @@ G_net = build_generator(n_input=1, n_output=2, size=256)
 G_net.load_state_dict(torch.load(model_path, map_location=device))
 G_net.eval()
 
-# Preprocessing function
 def preprocess_image(img):
     img = img.convert("RGB")
     img = transforms.Resize((256, 256), Image.BICUBIC)(img)
@@ -39,7 +37,6 @@ def preprocess_image(img):
     L = img_to_lab[[0], ...] / 50. - 1.
     return L.unsqueeze(0).to(device)
 
-# Inference function
 def colorize_image(img, model):
     L = preprocess_image(img)
     with torch.no_grad():
@@ -53,14 +50,12 @@ def colorize_image(img, model):
         rgb_imgs.append(img_rgb)
     return np.stack(rgb_imgs, axis=0)
 
-# Gradio interface
 def colorize(img):
     colorized_images = colorize_image(img, G_net)
     colorized_image = colorized_images[0]
     return Image.fromarray((colorized_image * 255).astype(np.uint8))
 
-# Create the Gradio interface
-iface = gr.Interface(
+app = gr.Interface(
     fn=colorize,
     inputs=gr.Image(type="pil", label="Upload Grayscale Image"),
     outputs=gr.Image(type="pil", label="Colorized Image"),
@@ -69,5 +64,4 @@ iface = gr.Interface(
     allow_flagging="never"
 )
 
-# Launch the interface
-iface.launch()
+app.launch()