Update to transformers-based model.

app.py

@@ -1,67 +1,47 @@
-import gradio as gr
-
-import pathlib
 import base64
 import re
+import os
+import pathlib
+import random
 import time
 from io import BytesIO
 
+from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
+import gradio as gr
 import imgkit
-import os
 from PIL import Image
-from fastai.callback.core import Callback
-from fastai.learner import *
-from fastai.torch_core import TitledStr
-from torch import tensor, Tensor
-from torch.distributions import Transform
-import random
+import torch
+from transformers import GPT2LMHeadModel, GPT2TokenizerFast, pipeline
 
-# These utility functions need to be in main (or otherwise where created) because fastai loads from that module, see:
-# https://docs.fast.ai/learner.html#load_learner
-from transformers import GPT2TokenizerFast
 
-import torch
-from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
+temp = pathlib.PosixPath
+pathlib.PosixPath = pathlib.WindowsPath
 
 gpu = False
 
 AUTH_TOKEN = os.environ.get('AUTH_TOKEN')
-
+BASE_MODEL = "gpt2"
+MERGED_MODEL = "gpt2-magic-card"
 
 if gpu:
-    pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", custom_pipeline="stable_diffusion_mega", torch_dtype=torch.float16, revision="fp16", use_auth_token=AUTH_TOKEN)
-    scheduler = EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config)
-    pipeline.scheduler = scheduler
-    pipeline.to("cuda")
+    image_pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16,
+                                                       revision="fp16", use_auth_token=AUTH_TOKEN)
+    scheduler = EulerAncestralDiscreteScheduler.from_config(image_pipeline.scheduler.config)
+    image_pipeline.scheduler = scheduler
+    image_pipeline.to("cuda")
 else:
-    pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5",
-                                                 custom_pipeline="stable_diffusion_mega", use_auth_token=AUTH_TOKEN)
-    scheduler = EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config)
-    pipeline.scheduler = scheduler
+    image_pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5",  use_auth_token=AUTH_TOKEN)
+    scheduler = EulerAncestralDiscreteScheduler.from_config(image_pipeline.scheduler.config)
+    image_pipeline.scheduler = scheduler
 
 # Huggingface Spaces have 16GB RAM and 8 CPU cores
 # See https://huggingface.co/docs/hub/spaces-overview#hardware-resources
 
-pretrained_weights = 'gpt2'
-tokenizer = GPT2TokenizerFast.from_pretrained(pretrained_weights)
-
-
-def tokenize(text):
-    toks = tokenizer.tokenize(text)
-    return tensor(tokenizer.convert_tokens_to_ids(toks))
-
-
-class TransformersTokenizer(Transform):
-    def __init__(self, tokenizer): self.tokenizer = tokenizer
-
-    def encodes(self, x):
-        return x if isinstance(x, Tensor) else tokenize(x)
-
-    def decodes(self, x): return TitledStr(self.tokenizer.decode(x.cpu().numpy()))
-
-
-class DropOutput(Callback):
-    def after_pred(self): self.learn.pred = self.pred[0]
+model = GPT2LMHeadModel.from_pretrained(MERGED_MODEL)
+tokenizer = GPT2TokenizerFast.from_pretrained(BASE_MODEL)
+END_TOKEN = '###'
+eos_id = tokenizer.encode(END_TOKEN)
+text_pipeline = pipeline('text-generation', model=model, tokenizer=tokenizer)
 
 
 def gen_card_text(name):
@@ -70,15 +50,9 @@ def gen_card_text(name):
     else:
       prompt = f"Name: {name}\r\n"
     print(f'GENERATING CARD TEXT with prompt: {prompt}')
-    prompt_ids = tokenizer.encode(prompt)
-    if gpu:
-        inp = tensor(prompt_ids)[None].cuda()  # Use .cuda() for torch GPU
-    else:
-        inp = tensor(prompt_ids)[None]
-    preds = learner.model.generate(inp, max_length=512, num_beams=5, temperature=1.5, do_sample=True,
-                                   repetition_penalty=1.2)
-    result = tokenizer.decode(preds[0].cpu().numpy())
-    result = result.split('###')[0].replace(r'\r\n', '\n').replace('\r', '').replace(r'\r', '')
+    output = text_pipeline(prompt, max_length=512, num_return_sequences=1, num_beams=5, temperature=1.5, do_sample=True,
+                           repetition_penalty=1.2, eos_token_id=eos_id)
+    result = output[0]['generated_text'].split("###")[0].replace(r'\r\n', '\n').replace('\r', '').replace(r'\r', '')
     print(f'GENERATING CARD COMPLETE')
     print(result)
     if name == '':
@@ -87,10 +61,6 @@ def gen_card_text(name):
     return name, result
 
 
-# init only once
-learner = load_learner('./colab-data-test/export.pkl',
-                       cpu=not gpu)  # cpu=False uses GPU; make sure installed torch is GPU e.g. `pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116`
-
 pathlib.Path('card_data').mkdir(parents=True, exist_ok=True)
 pathlib.Path('card_images').mkdir(parents=True, exist_ok=True)
 pathlib.Path('card_html').mkdir(parents=True, exist_ok=True)
@@ -109,7 +79,7 @@ def run(name):
     prompt_template = f"fantasy illustration of a {card_type} {name}, by Greg Rutkowski"
     print(f"GENERATING IMAGE FOR {prompt_template}")
     # Regarding sizing see https://huggingface.co/blog/stable_diffusion#:~:text=When%20choosing%20image%20sizes%2C%20we%20advise%20the%20following%3A
-    images = pipeline.text2img(prompt_template, width=512, height=368, num_inference_steps=20).images
+    images = image_pipeline(prompt_template, width=512, height=368, num_inference_steps=20).images
     card_image = None
     for image in images:
         save_name = get_savename('card_images', name, 'png')

gpt2-magic-card/config.json

@@ -0,0 +1,39 @@
+{
+  "_name_or_path": "gpt2",
+  "activation_function": "gelu_new",
+  "architectures": [
+    "GPT2LMHeadModel"
+  ],
+  "attn_pdrop": 0.1,
+  "bos_token_id": 50256,
+  "embd_pdrop": 0.1,
+  "eos_token_id": 50256,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "gpt2",
+  "n_ctx": 1024,
+  "n_embd": 768,
+  "n_head": 12,
+  "n_inner": null,
+  "n_layer": 12,
+  "n_positions": 1024,
+  "reorder_and_upcast_attn": false,
+  "resid_pdrop": 0.1,
+  "scale_attn_by_inverse_layer_idx": false,
+  "scale_attn_weights": true,
+  "summary_activation": null,
+  "summary_first_dropout": 0.1,
+  "summary_proj_to_labels": true,
+  "summary_type": "cls_index",
+  "summary_use_proj": true,
+  "task_specific_params": {
+    "text-generation": {
+      "do_sample": true,
+      "max_length": 50
+    }
+  },
+  "torch_dtype": "float16",
+  "transformers_version": "4.32.1",
+  "use_cache": true,
+  "vocab_size": 50257
+}

gpt2-magic-card/generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 50256,
+  "eos_token_id": 50256,
+  "transformers_version": "4.32.1"
+}

colab-data-test/export.pkl → gpt2-magic-card/pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:81e1bb9267d4e38f1de3eb4f103ec49537640ce7536e3dfcbce8cd673a38fb85
-size 512717919
+oid sha256:c3b0f3c17995a67032a94e998a790f5ea1ae1dbfd9d231c8bc7625b6556cb7d1
+size 248927581

requirements.txt

@@ -1,6 +1,5 @@
 accelerate
 diffusers
-fastai
 ftfy
 gradio
 imgkit