Views
No views yet
1import pickle
2import random
3import numpy as np
4
5import os
6import wget
7from zipfile import ZipFile
8
9
10def download_model(force = False):
11 if force == True: print('downloading model file size is 108 MB so it may take some time to complete...')
12 try:
13 url = "https://huggingface.co/thefcraft/prompt-generator-stable-diffusion/resolve/main/models.pickle.zip"
14 if force == True:
15 with open("models.pickle.zip", 'w'): pass
16 wget.download(url, "models.pickle.zip")
17 if not os.path.exists('models.pickle.zip'): wget.download(url, "models.pickle.zip")
18 print('Download zip file now extracting model')
19 with ZipFile("models.pickle.zip", 'r') as zObject: zObject.extractall()
20 print('extracted model .. now all done')
21 return True
22 except:
23 if force == False: return download_model(force=True)
24 print('Something went wrong\ndownload model via link: `https://huggingface.co/thefcraft/prompt-generator-stable-diffusion/tree/main`')
25try: os.chdir(os.path.abspath(os.path.dirname(__file__)))
26except: pass
27if not os.path.exists('models.pickle'): download_model()
28
29with open('models.pickle', 'rb')as f:
30 models = pickle.load(f)
31
32LORA_TOKEN = ''#'<|>LORA_TOKEN<|>'
33# WEIGHT_TOKEN = '<|>WEIGHT_TOKEN<|>'
34NOT_SPLIT_TOKEN = '<|>NOT_SPLIT_TOKEN<|>'
35
36def sample_next(ctx:str,model,k):
37
38 ctx = ', '.join(ctx.split(', ')[-k:])
39 if model.get(ctx) is None:
40 return " "
41 possible_Chars = list(model[ctx].keys())
42 possible_values = list(model[ctx].values())
43
44 # print(possible_Chars)
45 # print(possible_values)
46
47 return np.random.choice(possible_Chars,p=possible_values)
48
49def generateText(model, minLen=100, size=5):
50 keys = list(model.keys())
51 starting_sent = random.choice(keys)
52 k = len(random.choice(keys).split(', '))
53
54 sentence = starting_sent
55 ctx = ', '.join(starting_sent.split(', ')[-k:])
56
57 while True:
58 next_prediction = sample_next(ctx,model,k)
59 sentence += f", {next_prediction}"
60 ctx = ', '.join(sentence.split(', ')[-k:])
61 # if sentence.count('\n')>size: break
62 if '\n' in sentence: break
63 sentence = sentence.replace(NOT_SPLIT_TOKEN, ', ')
64 # sentence = re.sub(WEIGHT_TOKEN.replace('|', '\|'), lambda match: f":{random.randint(0,2)}.{random.randint(0,9)}", sentence)
65 # sentence = sentence.replace(":0.0", ':0.1')
66 # return sentence
67
68 prompt = sentence.split('\n')[0]
69 if len(prompt)<minLen:
70 prompt = generateText(model, minLen, size=1)[0]
71
72 size = size-1
73 if size == 0: return [prompt]
74 output = []
75 for i in range(size+1):
76 prompt = generateText(model, minLen, size=1)[0]
77 output.append(prompt)
78
79 return output
80if __name__ == "__main__":
81 for model in models: # models = [(model, neg_model), (nsfw, neg_nsfw), (sfw, neg_sfw)]
82 text = generateText(model[0], minLen=300, size=5)
83 text_neg = generateText(model[1], minLen=300, size=5)
84
85 # print('\n'.join(text))
86 for i in range(len(text)):
87 print(text[i])
88 # print('negativePrompt:')
89 print(text_neg[i])
90 print('----------------------------------------------------------------')
91 print('********************************************************************************************************************************************************')
92
93