Spaces:
Runtime error
Runtime error
| import base64 | |
| import pathlib | |
| import re | |
| import time | |
| from io import BytesIO | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from PIL import Image, ImageChops, ImageDraw | |
| from fastai.callback.core import Callback | |
| from fastai.learner import * | |
| from fastai.torch_core import TitledStr | |
| from html2image import Html2Image | |
| # from min_dalle import MinDalle | |
| from torch import tensor, Tensor, float16, float32 | |
| from torch.distributions import Transform | |
| from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler | |
| # 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 os | |
| AUTH_TOKEN = os.environ.get('AUTH_TOKEN') | |
| # update requirements.txt with: | |
| # C:\Users\Grant\PycharmProjects\test_space\venv\Scripts\pip3.exe freeze > requirements.txt | |
| # 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] | |
| # initialize only once | |
| # Takes about 2 minutes (126 seconds) to generate an image in Huggingface spaces on CPU | |
| # NOTE as of 2022-11-13 min-dalle is broken, switch to using a stable diffusion model for images | |
| # model = MinDalle( | |
| # models_root='./pretrained', | |
| # dtype=float32, | |
| # device='cpu', | |
| # is_mega=True, | |
| # is_reusable=True | |
| # ) | |
| # Download pipeline, but overwrite scheduler | |
| # Consider DPMSolverMultistepScheduler once added to diffusers | |
| from diffusers import EulerAncestralDiscreteScheduler | |
| scheduler = EulerAncestralDiscreteScheduler.from_config("runwayml/stable-diffusion-v1-5", subfolder="scheduler", | |
| use_auth_token=AUTH_TOKEN) | |
| pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", custom_pipeline="stable_diffusion_mega", | |
| torch_dtype=torch.float32, | |
| scheduler=scheduler, use_auth_token=AUTH_TOKEN) | |
| # pipe.enable_attention_slicing() | |
| # pipeline.to("cuda") | |
| def gen_image(prompt): | |
| prompt = f"{prompt}, fantasy painting by Greg Rutkowski" | |
| # See https://huggingface.co/spaces/pootow/min-dalle/blob/main/app.py | |
| # Hugging Space faces seems to run out of memory if grads are not disabled | |
| # torch.set_grad_enabled(False) | |
| print(f'RUNNING gen_image with prompt: {prompt}') | |
| images = pipeline.text2img(prompt, width=256, height=256, num_inference_steps=20).images | |
| # images = model.generate_images( | |
| # text=prompt, | |
| # seed=-1, | |
| # grid_size=1, # grid size above 2 causes out of memory on 12 GB 3080Ti; grid size 2 gives 4 images | |
| # is_seamless=False, | |
| # temperature=1, | |
| # top_k=256, | |
| # supercondition_factor=16, | |
| # is_verbose=True | |
| # ) | |
| print('COMPLETED GENERATION') | |
| # images = images.to('cpu').numpy() | |
| # images = images.astype(np.uint8) | |
| # return Image.fromarray(images[0]) | |
| return images[0] | |
| gpu = False | |
| # init only once | |
| learner = load_learner('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` | |
| def parse_monster_description(name, text): | |
| match = re.search(r"Description: (.*)", text) | |
| if not match: | |
| return f"{name} is a monster." | |
| description = match.group(1) | |
| print(description.split('.')[0]) | |
| return description.split('.')[0] | |
| def gen_monster_text(name): | |
| prompt = f"Name: {name}\r\n" | |
| print(f'GENERATING MONSTER 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', '') | |
| print(f'GENERATING MONSTER COMPLETE') | |
| print(result) | |
| return result | |
| def extract_text_for_header(text, header): | |
| match = re.search(fr"{header}: (.*)", text) | |
| if match is None: | |
| return '' | |
| return match.group(1) | |
| def remove_section(html, html_class): | |
| match = re.search(f'<li class="{html_class}"([\w\W])*?li>', html) | |
| if match is not None: | |
| html = html.replace(match.group(0), '') | |
| return html | |
| def format_monster_card(monster_text, image_data): | |
| print('FORMATTING MONSTER TEXT') | |
| # see giffyglyph's monster maker https://giffyglyph.com/monstermaker/app/ | |
| # Different Formatting style examples and some json export formats | |
| card = pathlib.Path('monsterMakerTemplate.html').read_text() | |
| if not isinstance(image_data, (bytes, bytearray)): | |
| card = card.replace('{image_data}', f'{image_data}') | |
| else: | |
| card = card.replace('{image_data}', f'data:image/png;base64,{image_data.decode("utf-8")}') | |
| name = extract_text_for_header(monster_text, 'Name') | |
| card = card.replace('{name}', name) | |
| monster_type = extract_text_for_header(monster_text, 'Type') | |
| card = card.replace('{monster_type}', monster_type) | |
| armor_class = extract_text_for_header(monster_text, 'Armor Class') | |
| card = card.replace('{armor_class}', armor_class) | |
| hit_points = extract_text_for_header(monster_text, 'Hit Points') | |
| card = card.replace('{hit_points}', hit_points) | |
| speed = extract_text_for_header(monster_text, 'Speed') | |
| card = card.replace('{speed}', speed) | |
| str_stat = extract_text_for_header(monster_text, 'STR') | |
| card = card.replace('{str_stat}', str_stat) | |
| dex_stat = extract_text_for_header(monster_text, 'DEX') | |
| card = card.replace('{dex_stat}', dex_stat) | |
| con_stat = extract_text_for_header(monster_text, 'CON') | |
| card = card.replace('{con_stat}', con_stat) | |
| int_stat = extract_text_for_header(monster_text, 'INT') | |
| card = card.replace('{int_stat}', int_stat) | |
| wis_stat = extract_text_for_header(monster_text, 'WIS') | |
| card = card.replace('{wis_stat}', wis_stat) | |
| cha_stat = extract_text_for_header(monster_text, 'CHA') | |
| card = card.replace('{cha_stat}', cha_stat) | |
| saving_throws = extract_text_for_header(monster_text, 'Saving Throws') | |
| card = card.replace('{saving_throws}', saving_throws) | |
| if not saving_throws: | |
| card = remove_section(card, 'monster-saves') | |
| skills = extract_text_for_header(monster_text, 'Skills') | |
| card = card.replace('{skills}', skills) | |
| if not skills: | |
| card = remove_section(card, 'monster-skills') | |
| damage_vulnerabilities = extract_text_for_header(monster_text, 'Damage Vulnerabilities') | |
| card = card.replace('{damage_vulnerabilities}', damage_vulnerabilities) | |
| if not damage_vulnerabilities: | |
| card = remove_section(card, 'monster-vulnerabilities') | |
| damage_resistances = extract_text_for_header(monster_text, 'Damage Resistances') | |
| card = card.replace('{damage_resistances}', damage_resistances) | |
| if not damage_resistances: | |
| card = remove_section(card, 'monster-resistances') | |
| damage_immunities = extract_text_for_header(monster_text, 'Damage Immunities') | |
| card = card.replace('{damage_immunities}', damage_immunities) | |
| if not damage_immunities: | |
| card = remove_section(card, 'monster-immunities') | |
| condition_immunities = extract_text_for_header(monster_text, 'Condition Immunities') | |
| card = card.replace('{condition_immunities}', condition_immunities) | |
| if not condition_immunities: | |
| card = remove_section(card, 'monster-conditions') | |
| senses = extract_text_for_header(monster_text, 'Senses') | |
| card = card.replace('{senses}', senses) | |
| if not senses: | |
| card = remove_section(card, 'monster-senses') | |
| languages = extract_text_for_header(monster_text, 'Languages') | |
| card = card.replace('{languages}', languages) | |
| if not languages: | |
| card = remove_section(card, 'monster-languages') | |
| challenge = extract_text_for_header(monster_text, 'Challenge') | |
| card = card.replace('{challenge}', challenge) | |
| if not challenge: | |
| card = remove_section(card, 'monster-challenge') | |
| description = extract_text_for_header(monster_text, 'Description') | |
| card = card.replace('{description}', description) | |
| match = re.search(r"Passives:\n([\w\W]*)", monster_text) | |
| if match is None: | |
| passives = '' | |
| else: | |
| passives = match.group(1) | |
| p = passives.split(':') | |
| if len(p) > 1: | |
| p = ":".join(p) | |
| p = p.split('\n') | |
| passives_data = '' | |
| for x in p: | |
| x = x.split(':') | |
| if len(x) > 1: | |
| trait = x[0] | |
| if trait == "Passives": | |
| continue | |
| if 'Action' in trait: | |
| break | |
| detail = ":".join(x[1:]) | |
| passives_data += f'<div class="monster-trait"><p><span class="name">{trait}</span> <span class="detail">{detail}</span></p></div>' | |
| card = card.replace('{passives}', passives_data) | |
| else: | |
| card = card.replace('{passives}', f'<div class="monster-trait"><p>{passives}</p></div>') | |
| match = re.search(r"Actions:\n([\w\W]*)", monster_text) | |
| if match is None: | |
| actions = '' | |
| else: | |
| actions = match.group(1) | |
| a = actions.split(':') | |
| if len(a) > 1: | |
| a = ":".join(a) | |
| a = a.split('\n') | |
| actions_data = '' | |
| for x in a: | |
| x = x.split(':') | |
| if len(x) > 1: | |
| action = x[0] | |
| if action == "Actions": | |
| continue | |
| if 'Passive' in action: | |
| break | |
| detail = ":".join(x[1:]) | |
| actions_data += f'<div class="monster-action"><p><span class="name">{action}</span> <span class="detail">{detail}</span></p></div>' | |
| card = card.replace('{actions}', actions_data) | |
| else: | |
| card = card.replace('{actions}', f'<div class="monster-action"><p>{actions}</p></div>') | |
| # TODO: Legendary actions, reactions, make column count for format an option (1 or 2 column layout) | |
| card = card.replace('Melee or Ranged Weapon Attack:', '<i>Melee or Ranged Weapon Attack:</i>') | |
| card = card.replace('Melee Weapon Attack:', '<i>Melee Weapon Attack:</i>') | |
| card = card.replace('Ranged Weapon Attack:', '<i>Ranged Weapon Attack:</i>') | |
| card = card.replace('Hit:', '<i>Hit:</i>') | |
| print('FORMATTING MONSTER TEXT COMPLETE') | |
| return card | |
| def pil_to_base64(image): | |
| print('CONVERTING PIL IMAGE TO BASE64 STRING') | |
| buffered = BytesIO() | |
| image.save(buffered, format="PNG") | |
| img_str = base64.b64encode(buffered.getvalue()) | |
| print('CONVERTING PIL IMAGE TO BASE64 STRING COMPLETE') | |
| return img_str | |
| hti = Html2Image(output_path='rendered_cards') | |
| def trim(im, border): | |
| bg = Image.new(im.mode, im.size, border) | |
| diff = ImageChops.difference(im, bg) | |
| bbox = diff.getbbox() | |
| if bbox: | |
| return im.crop(bbox) | |
| def crop_background(image): | |
| white = (255, 255, 255) | |
| ImageDraw.floodfill(image, (image.size[0] - 1, 0), white, thresh=50) | |
| image = trim(image, white) | |
| return image | |
| def html_to_png(html): | |
| print('CONVERTING HTML CARD TO PNG IMAGE') | |
| paths = hti.screenshot(html_str=html, css_file="monstermaker.css", save_as="test.png", size=(800, 1440)) | |
| path = paths[0] | |
| print('OPENING IMAGE FROM FILE') | |
| img = Image.open(path).convert("RGB") | |
| print('CROPPING BACKGROUND') | |
| img = crop_background(img) | |
| print('CONVERTING HTML CARD TO PNG IMAGE COMPLETE') | |
| return img | |
| def run(name: str) -> (Image, str, Image, str): | |
| start = time.time() | |
| print(f'BEGINNING RUN FOR {name}') | |
| if not name: | |
| placeholder_image = Image.new(mode="RGB", size=(256, 256)) | |
| return placeholder_image, 'No name provided; enter a name and try again', placeholder_image, '' | |
| text = gen_monster_text(name) | |
| description = parse_monster_description(name, text) | |
| pil = gen_image(description) | |
| image_data = pil_to_base64(pil) | |
| card_html = format_monster_card(text, image_data) | |
| card_image = html_to_png(card_html) | |
| end = time.time() | |
| print(f'RUN COMPLETED IN {int(end - start)} seconds') | |
| return card_image, text, pil, card_html | |
| app_description = ( | |
| """ | |
| # Create your own D&D monster! | |
| Enter a name, click Submit, and wait for about 4 minutes to see the result. | |
| """).strip() | |
| input_box = gr.Textbox(label="Enter a monster name", placeholder="Jabberwock") | |
| output_monster_card = gr.Image(label="Monster Card", type='pil', value="examples/jabberwock_card.png") | |
| output_text_box = gr.Textbox(label="Monster Text", value=pathlib.Path("examples/jabberwock.txt").read_text('utf-8')) | |
| output_monster_image = gr.Image(label="Monster Image", type='pil', value="examples/jabberwock.png") | |
| output_monster_html = gr.HTML(label="Monster HTML", visible=False, show_label=False) | |
| x = gr.components.Textbox() | |
| iface = gr.Interface(title="MonsterGen", theme="default", description=app_description, fn=run, inputs=[input_box], | |
| outputs=[output_monster_card, output_text_box, output_monster_image, output_monster_html]) | |
| iface.launch() | |
| # TODO: Add examples, larger language model?, document process, log silences, "Passives" => "Traits", log timestamps | |
| # Fine tune dalle-mini? https://blog.paperspace.com/dalle-mini/ | |
| # API works, assuming query takes no longer than 30 seconds (504 gateway timeout) | |
| # Looks like API page improvements are in progress: https://github.com/gradio-app/gradio/issues/1325 | |
| # Example code below: | |
| # import requests | |
| # r = requests.post(url='https://hf.space/embed/gstaff/test_space/+/api/predict', json={"data": [""]}, | |
| # timeout=None) | |
| # print(r.json()) | |
| # Looks like Huggingface uses the queue push api, then polls for status: | |
| # fetch("https://hf.space/embed/gstaff/test_space/api/queue/push/", { | |
| # "headers": { | |
| # "accept": "*/*", | |
| # "accept-language": "en-US,en;q=0.9", | |
| # "content-type": "application/json", | |
| # "sec-ch-ua": "\".Not/A)Brand\";v=\"99\", \"Google Chrome\";v=\"103\", \"Chromium\";v=\"103\"", | |
| # "sec-ch-ua-mobile": "?0", | |
| # "sec-ch-ua-platform": "\"Windows\"", | |
| # "sec-fetch-dest": "empty", | |
| # "sec-fetch-mode": "cors", | |
| # "sec-fetch-site": "same-origin" | |
| # }, | |
| # "referrer": "https://hf.space/embed/gstaff/test_space/+?__theme=light", | |
| # "referrerPolicy": "strict-origin-when-cross-origin", | |
| # "body": "{\"fn_index\":0,\"data\":[\"Jabberwock\"],\"action\":\"predict\",\"session_hash\":\"v9ehgfho3p\"}", | |
| # "method": "POST", | |
| # "mode": "cors", | |
| # "credentials": "omit" | |
| # }); | |
| # fetch("https://hf.space/embed/gstaff/test_space/api/queue/status/", { | |
| # "headers": { | |
| # "accept": "*/*", | |
| # "accept-language": "en-US,en;q=0.9", | |
| # "content-type": "application/json", | |
| # "sec-ch-ua": "\".Not/A)Brand\";v=\"99\", \"Google Chrome\";v=\"103\", \"Chromium\";v=\"103\"", | |
| # "sec-ch-ua-mobile": "?0", | |
| # "sec-ch-ua-platform": "\"Windows\"", | |
| # "sec-fetch-dest": "empty", | |
| # "sec-fetch-mode": "cors", | |
| # "sec-fetch-site": "same-origin" | |
| # }, | |
| # "referrer": "https://hf.space/embed/gstaff/test_space/+?__theme=light", | |
| # "referrerPolicy": "strict-origin-when-cross-origin", | |
| # "body": "{\"hash\":\"09f5369a7a414169aa48948bad5fd93d\"}", | |
| # "method": "POST", | |
| # "mode": "cors", | |
| # "credentials": "omit" | |
| # }); |