Create app.py

app.py

@@ -0,0 +1,82 @@
+import gradio as gr
+from transformers import Qwen2ForCausalLM, AutoTokenizer
+import torch
+
+
+
+torch.set_grad_enabled(False)
+model = Qwen2ForCausalLM.from_pretrained("Thouph/tag2prompt-qwen2-0.5b-v0.1")
+model.generation_config.max_new_tokens = None
+"""
+Otherwise you will get this warning
+Both `max_new_tokens` (=2048) and `max_length`(=512) seem to have been set. `max_new_tokens` will take precedence. Please refer to the documentation for more information. (https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)
+"""
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B")
+
+def post(
+        input_text,
+        temperature,
+        top_p,
+        top_k,
+        output_num_words,
+):
+    global model, processor
+
+    prompt = f"{input_text}\n{output_num_words}\n\n"
+    inputs = tokenizer(
+        prompt,
+        padding="do_not_pad",
+        max_length=512,
+        truncation=True,
+        return_tensors="pt",
+    )
+
+    generate_ids = model.generate(
+        **inputs,
+        max_length=512,
+        do_sample=True,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k
+    )
+    generated_text = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    generated_text = generated_text[len(prompt):]
+    return generated_text
+
+def main():
+
+    with gr.Blocks() as iface:
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                text_input = gr.TextArea(label="Tags (in Underscore Format)",)
+
+                temperature = gr.Slider(maximum=1., value=0.5, minimum=0., label='Temperature')
+                top_p = gr.Slider(maximum=1., value=0.8, minimum=0.1, label='Top P')
+                top_k = gr.Slider(maximum=100, value=20, minimum=1, step=1, label='Top K')
+                output_num_words = gr.Slider(maximum=512, value=100, minimum=1, step=1, label='Output Num Words')
+            with gr.Column(scale=1):
+                with gr.Column():
+                    caption_output = gr.Textbox(lines=1, label="Output")
+                    caption_button = gr.Button(
+                        value="Run tag2prompt", interactive=True, variant="primary"
+                    )
+                    caption_button.click(
+                        post,
+                        [
+                            text_input,
+                            temperature,
+                            top_p,
+                            top_k,
+                            output_num_words
+                        ],
+                        [caption_output],
+                    )
+
+    iface.launch()
+
+
+
+
+if __name__ == "__main__":
+    main()