Update README.md

README.md

@@ -8,12 +8,35 @@ pipeline_tag: text-generation
 
 ### How to use
 ```
-import gradio as gr
 import pickle
 import random
 import numpy as np
 
-with open('models.pickle', 'rb')as f:    
+import os
+import wget
+from zipfile import ZipFile
+
+
+def download_model(force = False):
+    if force == True: print('downloading model file size is 108 MB so it may take some time to complete...')
+    try:
+        url = "https://huggingface.co/thefcraft/prompt-generator-stable-diffusion/resolve/main/models.pickle.zip"
+        if force == True: 
+            with open("models.pickle.zip", 'w'): pass
+            wget.download(url, "models.pickle.zip") 
+        if not os.path.exists('models.pickle.zip'): wget.download(url, "models.pickle.zip") 
+        print('Download zip file now extracting model')
+        with ZipFile("models.pickle.zip", 'r') as zObject: zObject.extractall()
+        print('extracted model .. now all done')
+        return True
+    except:
+        if force == False: return download_model(force=True)
+        print('Something went wrong\ndownload model via link: `https://huggingface.co/thefcraft/prompt-generator-stable-diffusion/tree/main`')
+try: os.chdir(os.path.abspath(os.path.dirname(__file__)))
+except: pass
+if not os.path.exists('models.pickle'): download_model()
+
+with open('models.pickle', 'rb')as f:
     models = pickle.load(f)
 
 LORA_TOKEN = ''#'<|>LORA_TOKEN<|>'
@@ -45,7 +68,6 @@ def generateText(model, minLen=100, size=5):
         next_prediction = sample_next(ctx,model,k)
         sentence += f", {next_prediction}"
         ctx = ', '.join(sentence.split(', ')[-k:])
-
         # if sentence.count('\n')>size: break
         if '\n' in sentence: break
     sentence = sentence.replace(NOT_SPLIT_TOKEN, ', ')
@@ -65,36 +87,18 @@ def generateText(model, minLen=100, size=5):
         output.append(prompt)
 
     return output
+if __name__ == "__main__":
+    for model in models: # models = [(model, neg_model), (nsfw, neg_nsfw), (sfw, neg_sfw)]
+        text = generateText(model[0], minLen=300, size=5)
+        text_neg = generateText(model[1], minLen=300, size=5)
 
-def sentence_builder(quantity, minLen, Type, negative):
-    if Type == "NSFW": idx=1
-    elif Type == "SFW": idx=2
-    else: idx=0
-    model = models[idx]
-    output = ""
-    for i in range(quantity):
-        prompt = generateText(model[0], minLen=minLen, size=1)[0]
-        output+=f"PROMPT:  {prompt}\n\n"
-        if negative: 
-            negative_prompt = generateText(model[1], minLen=minLen, size=5)[0]
-            output+=f"NEGATIVE PROMPT:  {negative_prompt}\n"
-        output+="----------------------------------------------------------------"
-        output+="\n\n\n"
-    
-    return output[:-3]
-
+        # print('\n'.join(text))
+        for i in range(len(text)):
+            print(text[i])
+            # print('negativePrompt:')
+            print(text_neg[i])
+            print('----------------------------------------------------------------')
+        print('********************************************************************************************************************************************************')
 
-ui = gr.Interface(
-    sentence_builder,
-    [
-        gr.Slider(1, 10, value=4, label="Count", info="Choose between 1 and 10", step=1),
-        gr.Slider(100, 1000, value=300, label="minLen", info="Choose between 100 and 1000", step=50),
-        gr.Radio(["NSFW", "SFW", "BOTH"], label="TYPE", info="NSFW stands for NOT SAFE FOR WORK, so choose any one you want?"),
-        gr.Checkbox(label="negitive Prompt", info="Do you want to generate negative prompt as well as prompt?"),
-    ],
-    "text"
-)
 
-if __name__ == "__main__":
-    ui.launch()
 ```