first commit

.gitignore

@@ -0,0 +1,2 @@
+/javascript
+/style.css

README.md

@@ -0,0 +1,37 @@
+# training-picker
+
+Adds a tab to the webui that allows the user to automatically extract keyframes from video, and manually extract 512x512 crops of those frames for use in model training.
+
+![image](https://user-images.githubusercontent.com/2313721/200236386-5fed34df-03e4-4ea6-a653-e1b60393afcd.png)
+
+## Installation:
+
+1. Install [AUTOMATIC1111's Stable Diffusion Webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui)
+2. Install [ffmpeg](https://ffmpeg.org/) for your operating system
+3. Clone this repository into the `extensions` folder inside the webui
+
+Make sure you don't already have `python-ffmpeg` installed globally, as the library this program uses is `ffmpeg-python`, and your installation will conflict with it.
+
+## Usage:
+
+### Creating an extracted frame set
+
+1. Drop videos you want to extract cropped frames from into the `training-picker/videos` folder
+2. Open up the Training Picker tab of the webui
+3. Select one of the videos you placed in the `training-picker/videos` folder from the dropdown on the left
+4. Click 'Extract Frames`
+5. After the extraction finishes, a new keyframe set for the video should be selectable in the dropdown on the right
+6. Select the keyframe set, and the frames will appear in the browser below
+
+Optionally, you can also just supply a large collection of individual images you would like to work with directly by placing them into a folder within `training-picker/extracted-frames`.
+
+### Cropping
+
+* Scroll up and down to increase or decrease the size of the crop brush
+* Ctrl+scroll to adjust the aspect ratio of the crop, and middle-click to reset the aspect ratio to 1:1
+* Shift+scroll to adjust the size / aspect ratio by smaller increments
+* Click to save a crop at the brush's position
+* Navigate between frames in the collection by clicking the navigation buttons, entering a number into the counter, or by using the arrow keys / AD
+* Select an outfill method to outfill the non-square area of a rectangular crop into a square shape
+* Click "Bulk process frames with chosen outfill method" to automatically process every image in the current frame set using the outfill method chosen, outputting to the directory under "Save crops to:"
+* Crops will be saved to `training-picker/cropped-frames` by default

install.py

@@ -0,0 +1,2 @@
+import launch
+launch.run_pip("install --upgrade transformers", "Requirment of Prompt-Maker")

scripts/prompt_generator.py

@@ -0,0 +1,93 @@
+import subprocess
+import platform
+import math
+import json
+import sys
+import os
+import re
+from pathlib import Path
+
+import gradio as gr
+import numpy as np
+from tqdm import tqdm
+from PIL import Image, ImageFilter
+import cv2
+
+from modules.ui import create_refresh_button, folder_symbol
+from modules.shared import opts, OptionInfo
+from modules import shared, paths, script_callbacks
+from transformers import GPT2Tokenizer, GPT2LMHeadModel
+
+
+def on_ui_tabs():
+    with gr.Blocks(analytics_enabled=False) as prompt_generator:
+
+        # structure
+        with gr.Column():
+            with gr.Row():
+                promptTxt = gr.Textbox(
+                    lines=2, elem_id="promptTxt", label="Start of the prompt")
+        with gr.Column():
+            with gr.Row():
+                temp_slider = gr.Slider(
+                    elem_id="temp_slider", label="Temperature", interactive=True, minimum=0, maximum=1, value=0.9)
+                max_length_slider = gr.Slider(
+                    elem_id="max_length_slider", label="Max Length", interactive=True, minimum=1, maximum=200, step=1, value=80)
+                top_k_slider = gr.Slider(
+                    elem_id="top_k_slider", label="Top K", value=8, minimum=1, maximum=20, interactive=True)
+        with gr.Column():
+            with gr.Row():
+                repetition_penalty_slider = gr.Slider(
+                    elem_id="repetition_penalty_slider", label="Repetition Penalty", value=1.2, minimum=0, maximum=10, interactive=True)
+                num_return_sequences_slider = gr.Slider(
+                    elem_id="num_return_sequences_slider", label="How Many To Generate", value=5, minimum=1, maximum=20, interactive=True)
+        with gr.Column():
+            with gr.Row():
+                generateButton = gr.Button(
+                    value="Generate", elem_id="generate_button")
+        with gr.Column():
+            Results = gr.Text(elem_id="Results_textBox", interactive=False)
+
+        # events
+        def generate_longer_prompt(prompt, temperature, top_k,
+                                   max_length, repetition_penalty, num_return_sequences):
+            try:
+                tokenizer = GPT2Tokenizer.from_pretrained('distilgpt2')
+                tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+                model = GPT2LMHeadModel.from_pretrained(
+                    'FredZhang7/distilgpt2-stable-diffusion-v2')
+            except Exception as e:
+                print(f"Exception encountered while attempting to install tokenizer")
+            try:
+                print(f"Generate new prompt from: \"{prompt}\"")
+                input_ids = tokenizer(prompt, return_tensors='pt').input_ids
+                output = model.generate(input_ids, do_sample=True, temperature=temperature,
+                                        top_k=top_k, max_length=max_length,
+                                        num_return_sequences=num_return_sequences,
+                                        repetition_penalty=repetition_penalty,
+                                        penalty_alpha=0.6, no_repeat_ngram_size=1, early_stopping=True)
+                print("Generation complete!")
+                tempString = ""
+                for i in range(len(output)):
+                    tempString += tokenizer.decode(output[i], skip_special_tokens=True) + "\n"
+                return tempString
+            except Exception as e:
+                print(f"Exception encountered while attempting to generate prompt: {e}")
+                return gr.update(), f"Error: {e}"
+        generateButton.click(fn=generate_longer_prompt, inputs=[
+                             promptTxt,temp_slider , top_k_slider, max_length_slider, 
+                             repetition_penalty_slider, num_return_sequences_slider], 
+                             outputs=[Results])
+    return (prompt_generator, "Prompt Generator", "Prompt Generator"),
+
+# def on_ui_settings():
+#     picker_path = Path(paths.script_path) / "training-picker"
+#     section = ('training-picker', "Training Picker")
+#     opts.add_option("training_picker_fixed_size", OptionInfo(512, "Fixed size to resize images to", section=section))
+#     opts.add_option("training_picker_videos_path", OptionInfo(str(picker_path / "videos"), "Path to read videos from", section=section))
+#     opts.add_option("training_picker_framesets_path", OptionInfo(str(picker_path / "extracted-frames"), "Path to store extracted frame sets in", section=section))
+#     opts.add_option("training_picker_default_output_path", OptionInfo(str(picker_path / "cropped-frames"), "Default cropped image output directory", section=section))
+
+
+# script_callbacks.on_ui_settings(on_ui_settings)
+script_callbacks.on_ui_tabs(on_ui_tabs)