Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python | |
| # coding: utf-8 | |
| import os | |
| import openai | |
| import gradio as gr | |
| import torch | |
| from diffusers import StableDiffusionPipeline | |
| from torch import autocast | |
| from contextlib import nullcontext | |
| #from PIL import Image | |
| #from torchvision import transforms | |
| openai.api_key = os.getenv('openaikey') | |
| authtoken = os.getenv('authtoken') | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| context = autocast if device == "cuda" else nullcontext | |
| dtype = torch.float16 if device == "cuda" else torch.float32 | |
| pipe = StableDiffusionPipeline.from_pretrained("stale2000/sd-dnditem", torch_dtype=dtype, use_auth_token=authtoken) | |
| pipe = pipe.to(device) | |
| disable_safety = True | |
| if disable_safety: | |
| def null_safety(images, **kwargs): | |
| return images, False | |
| pipe.safety_checker = null_safety | |
| def create_files(): | |
| directory = 'C:\\Users\\brcwa\\OneDrive\\Desktop\\destinyCaptures\\dnd\\fullcaptionsimple\\full' | |
| for filename2 in os.listdir(directory): | |
| if not filename2.endswith('txt'): | |
| continue | |
| f = os.path.join(directory, filename2) | |
| # checking if it is a file | |
| if os.path.isfile(f): | |
| text_file = open(f, "r") | |
| lines = text_file.read() | |
| print(lines.split(',')[1] + "," + lines.split(',')[1]) | |
| #create_files() | |
| def createGPTPrompt(item_type, description): | |
| return item_type.split(",")[0].split(" ")[-1] + " of " + description | |
| def convert_lines(lines): | |
| key_arr = [] | |
| key_hash = {} | |
| for line in lines: | |
| key = line.split(",")[0] | |
| val = line.split(",")[1] | |
| key_arr.append(key) | |
| key_hash[key] = val | |
| return key_arr, key_hash | |
| def predict(dropdown, style_dropdown, manual_gpt_replacement, manual_sd_prompt, n_samples, history=[]): | |
| # gpt3 | |
| sd_input = "" | |
| gpt_input = "" | |
| description = style_dropdown | |
| if manual_sd_prompt != '': | |
| gpt_input = manual_gpt_replacement | |
| else: | |
| gpt_input = "Describe the mechanics of a 5th Edition DnD item called '" + createGPTPrompt(dropdown, description) + "' :" | |
| if manual_sd_prompt != '': | |
| sd_input = manual_sd_prompt | |
| else: | |
| sd_input = "dnditem, " + dropdown + ", " + style_hashmap[style_dropdown] + ", circle inner background and white outerbackground" | |
| response = openai.Completion.create( | |
| model="text-davinci-003", | |
| prompt=gpt_input, | |
| temperature=0.9, | |
| max_tokens=200, | |
| top_p=1, | |
| frequency_penalty=0, | |
| presence_penalty=0.6) | |
| # tokenize the new input sentence | |
| responseText = response["choices"][0]["text"] | |
| history.append((sd_input, responseText)) | |
| #img generation | |
| scale = 5.5 | |
| #with autocast("cuda"): | |
| # images = pipe(n_samples*[prompt], guidance_scale=scale).images | |
| with context("cuda"): | |
| images = pipe(n_samples*[sd_input], guidance_scale=scale, num_inference_steps=40).images | |
| return history, history, images | |
| #inputText = gr.Textbox(placeholder="input query") | |
| manual_gpt_query = gr.Textbox(placeholder="Input any query here, to replace the gpt query builder entirely.") | |
| manual_sd_prompt = gr.Textbox(placeholder="Input any query here, to replace the gpt query builder entirely.") | |
| choiceArr = ["a pair of boots", "a cloak", "a pair of gloves", "a helmet", "a necklace", "a ring", "a robe", "a rod", "a shield", "a staff", "a sword", "a wand"] | |
| dropdown = gr.Dropdown(label= "Item Type", choices=choiceArr, value="a pair of boots") | |
| text_file = open("styles.txt", "r") | |
| lines = text_file.read().split('\n') | |
| dropdown_arr, style_hashmap = convert_lines(lines) | |
| style_dropdown = gr.Dropdown(label= "Item Ability and Style", choices=dropdown_arr, value="ultimate evil") | |
| output_img = gr.Gallery(label="Generated image") | |
| output_img.style(grid=2) | |
| step_slide = gr.Slider(1, 4, value=2, step=1), | |
| slide = gr.Slider(label="Number of Images Generated", minimum=1, maximum=4, value=2, step=1) | |
| gr.Interface(fn=predict, | |
| inputs=[dropdown, style_dropdown, manual_gpt_query,manual_sd_prompt,slide,'state'], | |
| outputs=["chatbot",'state', output_img]).launch() | |