Upload 6 files

.dockerignore

@@ -0,0 +1,2 @@
+_pycache__pycache_
+*.pyc

.gitignore

@@ -0,0 +1,2 @@
+_pycache_
+*.pyc

app.py

@@ -0,0 +1,19 @@
+from flask import Flask, request, send_file
+from model import load_model, generate
+from io import BytesIO
+
+app = Flask(__name__)
+model = load_model()
+
+def serve_img(pilImage):
+  imageIO = BytesIO()
+  pilImage.save(imageIO, 'JPEG', quality=70)
+  imageIO.seek(0)
+  return imageIO
+
+
+@app.route("/",methods=['GET'])
+def index():
+  image = generate(model)
+  g = serve_img(image)
+  return send_file(g, mimetype="image/jpeg")

model.py

@@ -0,0 +1,55 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+from PIL import Image
+
+device = "cpu"
+weights_path = "./weights/face_generator_v2.pth"
+IN_CHANNELS = 100
+
+class FaceGenerator(nn.Module):
+    def __init__(self, in_channels):
+      super(FaceGenerator, self).__init__()
+      self.main = nn.Sequential(
+      nn.ConvTranspose2d(in_channels, 1024, 4, 2,0, bias=False),
+      nn.BatchNorm2d(1024),
+      nn.ReLU(True), # [batch_size, 1024, 2, 2]
+
+      nn.ConvTranspose2d(1024, 512, 4, 2,1, bias=False),
+      nn.BatchNorm2d(512),
+      nn.ReLU(True), # [batch_size, 512, 7, 7]
+
+      nn.ConvTranspose2d(512, 256, 4, 2,1, bias=False),
+      nn.BatchNorm2d(256),
+      nn.ReLU(True), # [batch_size, 256, 14, 14]
+
+      nn.ConvTranspose2d(256, 128, 4, 2,1, bias=False),
+      nn.BatchNorm2d(128),
+      nn.ReLU(True), # [batch_size, 256, 28, 28]
+      
+      nn.ConvTranspose2d(128,3, 4, 2,1, bias=False),
+      nn.Sigmoid(), # [batch_size, 1, 32, 32]
+    )
+    def forward(self, x):
+      return self.main(x)
+    
+    
+def load_model():
+  model = FaceGenerator(IN_CHANNELS).to(device=device)
+  model.load_state_dict(torch.load(weights_path, map_location=torch.device(device)), strict=True)
+  model = model.eval()
+  
+  print("[!] Model Loaded..")
+  return model
+
+def generate(model):
+  noise = torch.randn((1,IN_CHANNELS, 1, 1)).to(device)
+  image_op = model(noise).squeeze()
+  image_op = image_op.permute(1,2,0).detach().cpu().numpy()
+  image_op = image_op * 255.0
+  image_op = image_op.astype(np.uint8)
+  image_op = Image.fromarray(image_op)
+  image_op = image_op.resize((256, 256),resample=Image.ADAPTIVE)
+  return image_op    
+  

requirements.txt

@@ -0,0 +1 @@
+Flask==3.0.2

weights/face_generator_v2.pth

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