Update app.py

app.py

@@ -1,55 +1,72 @@
 import gradio as gr
-from diffusers import StableDiffusionPipeline
+from huggingface_hub import hf_hub_download
 import torch
-from transformers import pipeline
-
-# Use an available colorization pipeline from Hugging Face
-colorizer = pipeline("image-to-image", model="HuggingFaceM4/DeOldify", device_map="auto")
-
-def colorize_image(input_image):
-    output = colorizer(input_image)
-    return output[0]['image']
-
-def reset():
-    return None, None
-
-with gr.Blocks() as demo:
-    gr.Markdown("# 🎨 Image Colorization App")
-
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(label="Upload your grayscale image", type="pil")
-            clear_button = gr.Button("🔄 Reset / Clear")
-            download_button = gr.File(label="Download Colorized Image")
-        with gr.Column():
-            output_image = gr.Image(label="Colorized Image")
-
-    colorize_btn = gr.Button("✨ Colorize Image")
-
-    colorize_btn.click(
-        colorize_image,
-        inputs=input_image,
-        outputs=output_image
-    )
-
-    clear_button.click(
-        reset,
-        inputs=[],
-        outputs=[input_image, output_image]
-    )
-
-    def prepare_download(image):
-        if image:
-            path = "colorized_output.png"
-            image.save(path)
-            return path
-        else:
-            return None
-
-    output_image.change(
-        prepare_download,
-        inputs=output_image,
-        outputs=download_button
-    )
-
-demo.launch()
+from PIL import Image
+from torchvision import transforms
+from skimage.color import rgb2lab, lab2rgb
+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)
+
+# Define the generator model (same architecture as used during training)
+from fastai.vision.learner import create_body
+from torchvision.models import resnet34
+from fastai.vision.models.unet import DynamicUnet
+
+def build_generator(n_input=1, n_output=2, size=256):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    backbone = create_body(resnet34(), pretrained=True, n_in=n_input, cut=-2)
+    G_net = DynamicUnet(backbone, n_output, (size, size)).to(device)
+    return G_net
+
+# Initialize and load the model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+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)
+    img = np.array(img)
+    img_to_lab = rgb2lab(img).astype("float32")
+    img_to_lab = transforms.ToTensor()(img_to_lab)
+    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():
+        ab = model(L)
+    L = (L + 1.) * 50.
+    ab = ab * 110.
+    Lab = torch.cat([L, ab], dim=1).permute(0, 2, 3, 1).cpu().numpy()
+    rgb_imgs = []
+    for img in Lab:
+        img_rgb = lab2rgb(img)
+        rgb_imgs.append(img_rgb)
+    return np.stack(rgb_imgs, axis=0)[0]  # Return the first (and only) image
+
+# Gradio interface function
+def colorization_function(image):
+    colorized_image = colorize_image(image, G_net)
+    return colorized_image
+
+# Gradio Interface Setup
+iface = gr.Interface(
+    fn=colorization_function,
+    inputs=gr.inputs.Image(type="pil", label="Upload Grayscale Image"),
+    outputs=gr.outputs.Image(type="numpy", label="Colorized Image"),
+    title="Image Colorization with GAN",
+    description="Upload a grayscale image, and the model will colorize it using AI."
+)
+
+# Launch the Gradio interface
+iface.launch()