Update app/main.py

app/main.py

@@ -1,84 +1,72 @@
-import io
-import uuid
-import os
+import torch
+from torch import nn
+from torchvision import transforms
 from PIL import Image
-from fastapi import FastAPI, UploadFile, File
-from fastapi.responses import FileResponse
-from huggingface_hub import from_pretrained_fastai
 import gradio as gr
-import torch
-import uvicorn
+import io
+import base64
+import requests
 
-# ==========================================================
-# 🔧 CONFIGURATION
-# ==========================================================
-MODEL_ID = "Hammad712/GAN-Colorization-Model"  # 👉 change this if you want another model
+# Download model from Hugging Face Hub
+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)
+# 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)
 
-# ==========================================================
-# 🚀 LOAD MODEL
-# ==========================================================
-print(f"Loading model: {MODEL_ID}")
-learn = from_pretrained_fastai(MODEL_ID)
-print("✅ Model loaded successfully!")
+# Load the model (generator)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+generator = torch.load(model_path, map_location=device)
+generator.eval()
 
-# ==========================================================
-# 🧠 Colorization Function
-# ==========================================================
-def colorize_image(image: Image.Image):
-    if image.mode != "RGB":
-        image = image.convert("RGB")
-    pred = learn.predict(image)[0]
-    return pred
+# Define transforms
+transform_gray = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.Grayscale(),
+    transforms.ToTensor()
+])
 
-# ==========================================================
-# 🌐 FASTAPI APP
-# ==========================================================
-app = FastAPI(title="Image Colorization API")
+transform_color = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.ToTensor()
+])
 
-@app.post("/colorize")
-async def colorize_endpoint(file: UploadFile = File(...)):
-    """
-    Upload a black & white image -> get colorized image
-    """
-    try:
-        img_bytes = await file.read()
-        image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+to_pil = transforms.ToPILImage()
 
-        colorized = colorize_image(image)
-        output_filename = f"{uuid.uuid4()}.png"
-        output_path = os.path.join(RESULT_DIR, output_filename)
-        colorized.save(output_path)
+# Colorization function
+def colorize_image(input_image):
+    img = transform_gray(input_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
 
-        return FileResponse(output_path, media_type="image/png")
-    except Exception as e:
-        return {"error": str(e)}
+# ---- 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"
+)
 
-# ==========================================================
-# 🎨 GRADIO INTERFACE
-# ==========================================================
-def gradio_interface(image):
-    return colorize_image(image)
+# ---- API Endpoint ----
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import StreamingResponse
 
-title = "🎨 FastAI / HuggingFace Image Colorizer"
-description = "Upload a black & white photo to get a colorized version."
+app = FastAPI(title="GAN Colorization API")
 
-iface = gr.Interface(
-    fn=gradio_interface,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Image(type="pil", label="Colorized Output"),
-    title=title,
-    description=description,
-)
+@app.post("/colorize")
+async def colorize_api(file: UploadFile = File(...)):
+    image = Image.open(io.BytesIO(await file.read())).convert("RGB")
+    colorized = colorize_image(image)
+    buf = io.BytesIO()
+    colorized.save(buf, format="PNG")
+    buf.seek(0)
+    return StreamingResponse(buf, media_type="image/png")
 
-gradio_app = gr.mount_gradio_app(app, iface, path="/")
+# Mount Gradio UI on /
+import gradio as gr
+app = gr.mount_gradio_app(app, iface, path="/")
 
-# ==========================================================
-# ▶️ RUN LOCALLY OR IN HUGGINGFACE SPACE
-# ==========================================================
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)