app.py

#!/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=150,
    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 = ["none", "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)



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)

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()