More stuff

.gitignore

@@ -0,0 +1,2 @@
+__pycache__
+*.zip

app.py

@@ -5,34 +5,49 @@ from random import sample
 
 with open("meta.json") as f:
     meta = json.load(f)
+
+meta_s = [x for x in meta if x["rating"] == "s"]
+meta_g = [x for x in meta if x["rating"] == "g"]
 with open("tag2cat.json") as f:
     tag2cat = json.load(f)
-f = lambda cat: [
-    ", ".join(
-        t
-        for t in (tag.replace("_", " ") for tag in e["tag_string_general"].split())
-        if tag2cat.get(t, "") == cat
-    )
-    for e in meta
-]
-l = [f(c) for c in "Character Outfit Action Scene".split()]
+f = lambda cat, e: ", ".join(
+    t
+    for t in (tag.replace("_", " ") for tag in e["tag_string_general"].split())
+    if tag2cat.get(t, "") == cat
+)
+
+cats = "Character Outfit Action Scene".split()
 
 
-def do():
-    idxs = sample(range(len(meta)), 4)
-    parts = [ll[idx] for ll, idx in zip(l, idxs)]
+def do(ratings):
+    if not ratings:
+        raise gr.Error("Please select at least one of the ratings")
+    m = (
+        meta_g
+        if ratings[0] == "General"
+        else meta_s
+        if ratings[0] == "Sensitive"
+        else meta
+    )
+    idxs = sample(range(len(m)), 4)
+    parts = [f(cat, m[idx]) for cat, idx in zip(cats, idxs)]
     return (",\n".join(parts),) + tuple(
         _
         for p, idx in zip(parts, idxs)
         for _ in (
-            f"{p} ([danbooru post](https://danbooru.donmai.us/posts/{meta[idx]['id']}))",
-            f'![]({meta[idx]["large_file_url"]})',
+            f"{p} ([danbooru post](https://danbooru.donmai.us/posts/{m[idx]['id']}))",
+            f'![]({m[idx]["large_file_url"]})',
         )
     )
 
 
 with gr.Blocks() as demo:
     with gr.Row():
+        ratings = gr.CheckboxGroup(
+            ["General", "Sensitive"],
+            label="Allowed ratings",
+            value=["General", "Sensitive"],
+        )
         button = gr.Button("Get Random Prompt", variant="primary")
     with gr.Row():
         joined = gr.Text(label="generated prompt", show_copy_button=True)
@@ -48,7 +63,7 @@ with gr.Blocks() as demo:
         img_scene = gr.Markdown()
     button.click(
         do,
-        [],
+        [ratings],
         [
             joined,
             text_char,