chore: use gpu

app.py

@@ -1,6 +1,6 @@
 import os
 import time
-
+import spaces
 
 import torch
 from transformers import (
@@ -16,7 +16,7 @@ import gradio as gr
 MODEL_NAME = os.environ.get("MODEL_NAME", None)
 assert MODEL_NAME is not None
 MODEL_PATH = hf_hub_download(repo_id=MODEL_NAME, filename="model.safetensors")
-DEVICE = torch.device("cpu")
+DEVICE = torch.device("cuda")
 
 
 def fix_compiled_state_dict(state_dict: dict):
@@ -24,12 +24,11 @@ def fix_compiled_state_dict(state_dict: dict):
 
 
 def prepare_models():
-    config = AutoConfig.from_pretrained(
-        MODEL_NAME, use_cache=True, trust_remote_code=True
-    )
+    config = AutoConfig.from_pretrained(MODEL_NAME, trust_remote_code=True)
     model = AutoModelForPreTraining.from_config(
         config, torch_dtype=torch.bfloat16, trust_remote_code=True
     )
+    model.decoder_model.use_cache = True
     processor = AutoProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True)
 
     state_dict = load_file(MODEL_PATH)
@@ -38,7 +37,7 @@ def prepare_models():
 
     model.eval()
     model = model.to(DEVICE)
-    model = torch.compile(model)
+    # model = torch.compile(model)
 
     return model, processor
 
@@ -46,6 +45,7 @@ def prepare_models():
 def demo():
     model, processor = prepare_models()
 
+    @spaces.GPU(duration=5)
     @torch.inference_mode()
     def generate_tags(
         text: str,
@@ -109,8 +109,8 @@ def demo():
                         label="Auto detect copyright tags.", value=False
                     )
                     copyright_tags = gr.Textbox(
-                        label="Custom tags",
-                        placeholder="Enter custom tags here. e.g.) hatsune miku",
+                        label="Copyright tags",
+                        placeholder="Enter copyright tags here. e.g.) hatsune miku",
                     )
                     translate_btn = gr.Button(value="Translate")
 
@@ -124,9 +124,12 @@ def demo():
                             value=0.1,
                             step=0.1,
                         )
-                        top_k = gr.Number(
+                        top_k = gr.Slider(
                             label="Top k",
+                            minimum=1,
+                            maximum=100,
                             value=10,
+                            step=1,
                         )
                         top_p = gr.Slider(
                             label="Top p",

requirements.txt

@@ -3,3 +3,4 @@ transformers
 accelerate
 safetensors
 huggingface_hub
+spaces