feat: added app

app.py

@@ -0,0 +1,27 @@
+import gradio as gr
+
+from service import Service
+
+prompt = ""
+scale = 10
+n_samples = 4
+rows = 2
+cols = 2
+
+service = Service()
+
+
+def generate_pkm(prompt):
+    return service.generate(prompt, n_samples=n_samples, rows=rows, cols=cols)
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("Write a prompt to generate a Pokemon!")
+    prompt = gr.Textbox(label="Prompt")
+    image_output = gr.Image()
+    generate = gr.Button("Generate")
+    generate.click(generate_pkm, inputs=prompt, outputs=image_output)
+gr.Markdown("Models finetuned by: lambdalabs")
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,7 @@
+gradio
+diffusers==0.8.0
+transformers #==4.24.0
+scipy
+ftfy
+torch
+accelerate

service.py

@@ -0,0 +1,30 @@
+import torch
+from PIL import Image
+from diffusers import StableDiffusionPipeline
+from torch import autocast
+
+pipe = StableDiffusionPipeline.from_pretrained("lambdalabs/sd-pokemon-diffusers", torch_dtype=torch.float16)
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+pipe = pipe.to(device)
+
+
+class Service():
+    def __init__(self):
+        self.scale = 10
+
+    def image_grid(imgs, rows, cols):
+        assert len(imgs) == rows * cols
+        w, h = imgs[0].size
+        grid = Image.new('RGB', size=(cols * w, rows * h))
+        for i, img in enumerate(imgs):
+            grid.paste(img, box=(i % cols * w, i // cols * h))
+        return grid
+
+    def generate(self,
+                 prompt,
+                 n_samples,
+                 rows,
+                 cols):
+        with autocast("device"):
+            images = pipe(n_samples * [prompt], guidance_scale=self.scale).images
+        return self.image_grid(images, rows=rows, cols=cols)