chore: run on cpu

app.py

@@ -16,6 +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")
 
 
 def fix_compiled_state_dict(state_dict: dict):
@@ -36,6 +37,7 @@ def prepare_models():
     model.load_state_dict(state_dict)
 
     model.eval()
+    model = model.to(DEVICE)
     model = torch.compile(model)
 
     return model, processor
@@ -48,7 +50,7 @@ def demo():
     def generate_tags(
         text: str,
         auto_detect: bool,
-        copyright_tags: str,
+        copyright_tags: str = "",
         max_new_tokens: int = 128,
         do_sample: bool = False,
         temperature: float = 0.1,
@@ -70,10 +72,10 @@ def demo():
 
         start_time = time.time()
         outputs = model.generate(
-            input_ids=inputs["input_ids"].to("cuda"),
-            attention_mask=inputs["attention_mask"].to("cuda"),
-            encoder_input_ids=inputs["encoder_input_ids"].to("cuda"),
-            encoder_attention_mask=inputs["encoder_attention_mask"].to("cuda"),
+            input_ids=inputs["input_ids"].to(model.device),
+            attention_mask=inputs["attention_mask"].to(model.device),
+            encoder_input_ids=inputs["encoder_input_ids"].to(model.device),
+            encoder_attention_mask=inputs["encoder_attention_mask"].to(model.device),
             max_new_tokens=max_new_tokens,
             do_sample=do_sample,
             temperature=temperature,
@@ -93,44 +95,50 @@ def demo():
         )
         return [deocded, f"Time elapsed: {elapsed:.2f} seconds"]
 
+    # warmup
+    print("warming up...")
+    print(generate_tags("Miku is looking at viewer.", True))
+    print("done.")
+
     with gr.Blocks() as ui:
-        with gr.Row():
-            with gr.Column():
-                text = gr.Text(label="Text", lines=4)
-                auto_detect = gr.Checkbox(
-                    label="Auto detect copyright tags.", value=False
-                )
-                copyright_tags = gr.Textbox(
-                    label="Custom tags",
-                    placeholder="Enter custom tags here. e.g.) hatsune miku",
-                )
-                translate_btn = gr.Button(value="Translate")
-
-                with gr.Accordion(label="Advanced", open=False):
-                    max_new_tokens = gr.Number(label="Max new tokens", value=128)
-                    do_sample = gr.Checkbox(label="Do sample", value=False)
-                    temperature = gr.Slider(
-                        label="Temperature",
-                        minimum=0.1,
-                        maximum=1.0,
-                        value=0.1,
-                        step=0.1,
-                    )
-                    top_k = gr.Number(
-                        label="Top k",
-                        value=10,
+        with gr.Column():
+            with gr.Row():
+                with gr.Column():
+                    text = gr.Text(label="Text", lines=4)
+                    auto_detect = gr.Checkbox(
+                        label="Auto detect copyright tags.", value=False
                     )
-                    top_p = gr.Slider(
-                        label="Top p",
-                        minimum=0.1,
-                        maximum=1.0,
-                        value=0.1,
-                        step=0.1,
+                    copyright_tags = gr.Textbox(
+                        label="Custom tags",
+                        placeholder="Enter custom tags here. e.g.) hatsune miku",
                     )
-
-            with gr.Column():
-                output = gr.Textbox(label="Output", lines=4, interactive=False)
-                time_elapsed = gr.Markdown(value="")
+                    translate_btn = gr.Button(value="Translate")
+
+                    with gr.Accordion(label="Advanced", open=False):
+                        max_new_tokens = gr.Number(label="Max new tokens", value=128)
+                        do_sample = gr.Checkbox(label="Do sample", value=False)
+                        temperature = gr.Slider(
+                            label="Temperature",
+                            minimum=0.1,
+                            maximum=1.0,
+                            value=0.1,
+                            step=0.1,
+                        )
+                        top_k = gr.Number(
+                            label="Top k",
+                            value=10,
+                        )
+                        top_p = gr.Slider(
+                            label="Top p",
+                            minimum=0.1,
+                            maximum=1.0,
+                            value=0.1,
+                            step=0.1,
+                        )
+
+                with gr.Column():
+                    output = gr.Textbox(label="Output", lines=4, interactive=False)
+                    time_elapsed = gr.Markdown(value="")
 
             gr.Examples(
                 examples=[["Miku is looking at viewer.", True]],
@@ -139,9 +147,9 @@ def demo():
 
         gr.on(
             triggers=[
-                text.change,
-                auto_detect.change,
-                copyright_tags.change,
+                # text.change,
+                # auto_detect.change,
+                # copyright_tags.change,
                 translate_btn.click,
             ],
             fn=generate_tags,