add face gen

imagegen.py

@@ -0,0 +1,34 @@
+from PIL import Image
+import numpy as np
+from fastapi.responses import StreamingResponse
+from io import BytesIO
+from matplotlib import pyplot
+import matplotlib
+
+class FaceGen:
+    def __init__(self,model,device):
+        self.model = model
+        self.device = device
+
+    def generate_latent_points(self,latent_dim,n_samples,seed):
+        z_input = seed.reshape(n_samples,latent_dim)
+        return z_input
+
+    def plot_generated(self, examples, n):
+        buf = BytesIO()
+        for i in range(n * n):
+            pyplot.subplot(n, n, 1 + i)
+            pyplot.axis('off')
+            pyplot.imshow(examples[i, :, :])
+        pyplot.savefig(buf, format='png', transparent=True)
+        buf.seek(0)
+        return StreamingResponse(buf, media_type="image/png")
+
+    def generate_image(self, latent_dim, n_samples, seed):
+        latent_points = self.generate_latent_points(latent_dim, n_samples, seed)
+        generated_images = self.model.predict(latent_points)
+        X = (generated_images + 1) / 2.0
+        return self.plot_generated(X, 1)
+
+
+    

main.py

@@ -4,11 +4,19 @@ from datasets import load_dataset
 import pandas as pd
 import re
 
+from numpy.random import randn
+from keras.models import load_model
+from matplotlib import pyplot
+import matplotlib
+import tensorflow as tf
+
+
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 
 import chatbot
 import textgen
+import imagegen
 
 
 app = FastAPI()
@@ -87,4 +95,28 @@ async def generate_text(input: TextGenInput):
 
     response = model.generate_text(input.user)
 
+    return response
+
+@app.get("/face-gen")
+async def generate_face():
+    print("Generating face image request")
+
+    global model
+    
+    try:
+        # check if model is already loaded
+        if not isinstance(model, imagegen.FaceGen):
+            model = imagegen.FaceGen(
+                model = load_model('models/face-gen-gan/generator_model_100.h5'),
+                device =  '/CPU:0'
+            )
+    except NameError:
+        # if model is not defined, load
+        model = imagegen.FaceGen(
+            model = load_model('models/face-gen-gan/generator_model_100.h5'),
+            device = '/GPU:0' if tf.config.list_physical_devices('GPU') else '/CPU:0'
+        )
+
+    response = model.generate_image(latent_dim=100,n_samples=128,seed=randn(100*128))
+
     return response

models/face-gen-gan/generator_model_100.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd167921ef33931acace69054326314197b8bc340bd702503f167710f21f7235
+size 5563784

requirements.txt

@@ -4,4 +4,9 @@ torch==2.3.1
 datasets==2.21.0
 pandas==2.1.4
 uvicorn[standard]
-fastapi[all]
+fastapi[all]
+numpy==1.26.4
+matplotlib==3.9.1
+tensorflow==2.15.0
+keras==2.15.0
+pillow