Update app/main.py

app/main.py

@@ -1,72 +1,112 @@
+import io
+import os
+import uuid
 import torch
-from torch import nn
-from torchvision import transforms
+import torch.nn as nn
+from fastapi import FastAPI, UploadFile, File
+from fastapi.responses import FileResponse
 from PIL import Image
+import torchvision.transforms as T
 import gradio as gr
-import io
-import base64
-import requests
+import uvicorn
 
-# Download model from Hugging Face Hub
+# ==========================================================
+# 🔧 PATHS
+# ==========================================================
 MODEL_PATH = "generator.pt"
-HF_MODEL_REPO = "Hammad712/GAN-Colorization-Model"
+UPLOAD_DIR = "/tmp/uploads"
+RESULT_DIR = "/tmp/results"
+os.makedirs(UPLOAD_DIR, exist_ok=True)
+os.makedirs(RESULT_DIR, exist_ok=True)
+
+# ==========================================================
+# 🧩 Define Generator Architecture (from repo style)
+# ==========================================================
+class Generator(nn.Module):
+    def __init__(self):
+        super(Generator, self).__init__()
+        self.main = nn.Sequential(
+            nn.Conv2d(1, 64, kernel_size=4, stride=2, padding=1),
+            nn.ReLU(True),
+
+            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(True),
+
+            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(True),
 
-# Auto-download model if not present
-from huggingface_hub import hf_hub_download
-model_path = hf_hub_download(repo_id=HF_MODEL_REPO, filename=MODEL_PATH)
+            nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1),
+            nn.Tanh()
+        )
 
-# Load the model (generator)
+    def forward(self, x):
+        return self.main(x)
+
+# ==========================================================
+# 🚀 Load Model
+# ==========================================================
 device = "cuda" if torch.cuda.is_available() else "cpu"
-generator = torch.load(model_path, map_location=device)
-generator.eval()
+generator = Generator().to(device)
 
-# Define transforms
-transform_gray = transforms.Compose([
-    transforms.Resize((256, 256)),
-    transforms.Grayscale(),
-    transforms.ToTensor()
-])
+# Load weights
+state_dict = torch.load(MODEL_PATH, map_location=device)
+generator.load_state_dict(state_dict)
+generator.eval()
 
-transform_color = transforms.Compose([
-    transforms.Resize((256, 256)),
-    transforms.ToTensor()
-])
+print("✅ Model loaded successfully!")
 
-to_pil = transforms.ToPILImage()
+# ==========================================================
+# 🎨 Colorization Function
+# ==========================================================
+def colorize_image(image: Image.Image):
+    transform = T.Compose([
+        T.Resize((256, 256)),
+        T.Grayscale(num_output_channels=1),
+        T.ToTensor()
+    ])
 
-# Colorization function
-def colorize_image(input_image):
-    img = transform_gray(input_image).unsqueeze(0).to(device)
+    img_tensor = transform(image).unsqueeze(0).to(device)
     with torch.no_grad():
-        output = generator(img)
-    output_image = to_pil(output.squeeze().cpu().clamp(0, 1))
-    return output_image
-
-# ---- Gradio Interface ----
-iface = gr.Interface(
-    fn=colorize_image,
-    inputs=gr.Image(type="pil", label="Upload Grayscale Image"),
-    outputs=gr.Image(type="pil", label="Colorized Image"),
-    title="GAN Image Colorization (Hammad712)",
-    description="Colorizes black and white images using GAN model"
-)
+        output = generator(img_tensor)
+    output = (output.squeeze(0).permute(1, 2, 0).cpu().numpy() + 1) / 2.0  # Scale 0-1
 
-# ---- API Endpoint ----
-from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import StreamingResponse
+    output_img = Image.fromarray((output * 255).astype("uint8"))
+    return output_img
 
-app = FastAPI(title="GAN Colorization API")
+# ==========================================================
+# 🌐 FASTAPI APP
+# ==========================================================
+app = FastAPI(title="GAN Image Colorization API")
 
 @app.post("/colorize")
-async def colorize_api(file: UploadFile = File(...)):
-    image = Image.open(io.BytesIO(await file.read())).convert("RGB")
+async def colorize_endpoint(file: UploadFile = File(...)):
+    img_bytes = await file.read()
+    image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
     colorized = colorize_image(image)
-    buf = io.BytesIO()
-    colorized.save(buf, format="PNG")
-    buf.seek(0)
-    return StreamingResponse(buf, media_type="image/png")
 
-# Mount Gradio UI on /
-import gradio as gr
-app = gr.mount_gradio_app(app, iface, path="/")
+    output_filename = f"{uuid.uuid4()}.png"
+    output_path = os.path.join(RESULT_DIR, output_filename)
+    colorized.save(output_path)
+
+    return FileResponse(output_path, media_type="image/png")
+
+# ==========================================================
+# 💠 GRADIO UI
+# ==========================================================
+def gradio_ui(image):
+    return colorize_image(image)
+
+iface = gr.Interface(
+    fn=gradio_ui,
+    inputs=gr.Image(type="pil", label="Upload B&W Image"),
+    outputs=gr.Image(type="pil", label="Colorized Image"),
+    title="🎨 GAN Image Colorization",
+    description="Upload a black-and-white photo to get it colorized using a GAN model."
+)
+
+gradio_app = gr.mount_gradio_app(app, iface, path="/")
 
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)