Create thumbnail_generator.py

thumbnail_generator.py

@@ -0,0 +1,15 @@
+import torch
+from torchvision.utils import save_image
+from PIL import Image
+import io
+
+class ThumbnailGenerator:
+    def __init__(self, model_path="pytorch_model.bin"):
+        self.model = torch.load(model_path, map_location=torch.device("cpu"))
+        self.model.eval()
+
+    def generate(self, seed=None):
+        z = torch.randn(1, 512) if seed is None else torch.tensor(seed).float().unsqueeze(0)
+        with torch.no_grad():
+            thumbnail = self.model(z)
+        return thumbnail.squeeze(0)