Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import random
 from time import time, ctime
 
-def promptgen(choice, num, artist):
+def promptgen2(choice, num, artist):
   t = time()
   print(ctime(t))
 
@@ -30,6 +30,40 @@ def promptgen(choice, num, artist):
   print(', '.join(set(generated)) + '\n')
   return ', '.join(set(generated))
 
+def promptgen(choice, num, artist):
+  try:
+      gen_prompt = open('gen_prompt.txt').read().splitlines()
+  except:
+      pass
+  t = time()
+  print(ctime(t))
+
+  if choice == "Prompt Generator v0.1(Better quality)":
+    prompt = open('pr1.txt').read().splitlines()
+  elif choice == "Prompt Generator v0.2(More tags)":
+    prompt = open('pr2.txt').read().splitlines()
+
+  if int(num) < 1 or int(num) > 20:
+    num = 10
+  
+  if int(artist) < 0 or int(artist) > 40:
+    artist = 2
+
+  vocab = len(prompt)
+  generated = []
+  artists_num = 0
+  with open('gen_prompt.txt', 'wb') as f:
+      while len(sorted(set(generated), key=lambda d: generated.index(d))) < num:
+        rand = random.randint(0, vocab)
+        if prompt[rand-1].startswith('art by') and artists_num < artist:
+            artists_num +=1 
+            generated.append(prompt[rand-1])
+        elif not prompt[rand-1].startswith('art by'):
+            generated.append(prompt[rand-1])
+      print(', '.join(set(generated)) + '\n')
+      f.write(f'{ctime(t)} - {', '.join(set(generated))}')
+      return ', '.join(set(generated))
+
 demo = gr.Blocks()
 
 with demo: