Fix model path to point to .h5 file

Dockerfile

@@ -16,7 +16,7 @@ RUN pip install --upgrade pip
 COPY requirements.txt .
 
 # Install packages separately to handle large downloads better
-RUN pip install --default-timeout=1000 --no-cache-dir fastapi[all] uvicorn[standard] pydantic
+RUN pip install --default-timeout=1000 --no-cache-dir flask flask-cors
 
 # Install TensorFlow CPU (smaller and faster to download than GPU version)
 RUN pip install --default-timeout=1000 --no-cache-dir tensorflow-cpu==2.15.0
@@ -34,4 +34,4 @@ COPY models/ ./models/
 EXPOSE 8002
 
 # Run the application
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8002"]
+CMD ["python", "app.py"]

app.py

@@ -1,78 +1,70 @@
-from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import StreamingResponse
-from pydantic import BaseModel
+from flask import Flask, jsonify, request, send_file
+from flask_cors import CORS
 import numpy as np
 from keras.models import load_model
 from PIL import Image
 import io
-from typing import Optional
 
-app = FastAPI(title="Face Generator API", version="1.0.0")
+app = Flask(__name__)
 
-# Add CORS middleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+# Enable CORS for all routes
+CORS(app)
 
 # Global variables for model
-MODEL_PATH = "./models/face-gen-gan"
+MODEL_PATH = "./models/face-gen-gan/generator_model_100.h5"
 model = None
 latent_dim = None
 
-# Request models
-class FaceGenRequest(BaseModel):
-    n_samples: Optional[int] = 1
-    seed: Optional[int] = None
-
-@app.on_event("startup")
-async def load_gan_model():
-    """Load the GAN model on startup"""
+def load_gan_model():
+    """Load the GAN model"""
     global model, latent_dim
-    print(f"Loading face generation GAN model from {MODEL_PATH}...")
-    
-    model = load_model(MODEL_PATH)
-    latent_dim = model.input_shape[1]
-    
-    print(f"Model loaded successfully! Latent dimension: {latent_dim}")
+    if model is None:
+        print(f"Loading face generation GAN model from {MODEL_PATH}...")
+        model = load_model(MODEL_PATH)
+        latent_dim = model.input_shape[1]
+        print(f"Model loaded successfully! Latent dimension: {latent_dim}")
 
-@app.get("/")
+# Load model on startup
+load_gan_model()
+
+@app.route("/")
 def root():
-    return {
+    return jsonify({
         "message": "Face Generator API",
         "status": "running",
         "model": "face-gen-gan",
         "latent_dim": latent_dim
-    }
+    })
 
-@app.get("/health")
+@app.route("/health")
 def health():
-    return {
+    return jsonify({
         "status": "healthy",
         "model_loaded": model is not None,
         "latent_dim": latent_dim
-    }
+    })
 
-@app.post("/generate")
-async def generate_faces(request: FaceGenRequest):
+@app.route("/generate", methods=["POST"])
+def generate_faces():
     """
     Generate face images using the GAN model
     Returns a PNG image (single face or grid of faces)
     """
     if model is None:
-        return {"error": "Model not loaded"}
+        return jsonify({"error": "Model not loaded"}), 500
     
     try:
+        # Get request data
+        data = request.get_json() or {}
+        n_samples = data.get("n_samples", 1)
+        seed = data.get("seed", None)
+        
         # Validate n_samples
-        n_samples = max(1, min(request.n_samples, 16))  # Limit to 1-16
+        n_samples = max(1, min(int(n_samples), 16))  # Limit to 1-16
         
         # Set seed if provided
-        if request.seed is not None:
-            np.random.seed(request.seed)
+        if seed is not None:
+            np.random.seed(int(seed))
         
         # Generate random latent points
         latent_points = np.random.randn(n_samples, latent_dim)
@@ -108,19 +100,47 @@ async def generate_faces(request: FaceGenRequest):
         img.save(buf, format='PNG')
         buf.seek(0)
         
-        return StreamingResponse(buf, media_type="image/png")
+        return send_file(buf, mimetype='image/png')
     
     except Exception as e:
-        return {"error": str(e)}
+        return jsonify({"error": str(e)}), 500
 
-@app.get("/generate-single")
-async def generate_single_face(seed: Optional[int] = None):
+@app.route("/generate-single")
+def generate_single_face():
     """
     Quick endpoint to generate a single face
     """
-    request = FaceGenRequest(n_samples=1, seed=seed)
-    return await generate_faces(request)
+    seed = request.args.get('seed', None)
+    
+    if model is None:
+        return jsonify({"error": "Model not loaded"}), 500
+    
+    try:
+        # Set seed if provided
+        if seed is not None:
+            np.random.seed(int(seed))
+        
+        # Generate random latent points
+        latent_points = np.random.randn(1, latent_dim)
+        
+        # Generate images
+        generated_images = model.predict(latent_points, verbose=0)
+        
+        # Scale from [-1, 1] to [0, 255]
+        generated_images = ((generated_images + 1) / 2.0 * 255).astype(np.uint8)
+        
+        # Single image
+        img = Image.fromarray(generated_images[0])
+        
+        # Convert to bytes
+        buf = io.BytesIO()
+        img.save(buf, format='PNG')
+        buf.seek(0)
+        
+        return send_file(buf, mimetype='image/png')
+    
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
 
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8002)
+    app.run(host="0.0.0.0", port=8002, debug=False)

requirements.txt

@@ -1,7 +1,6 @@
-fastapi[all]
-uvicorn[standard]
+flask
+flask-cors
 tensorflow-cpu==2.15.0
 keras==2.15.0
 numpy==1.26.4
 pillow
-pydantic