Update app.py

app.py

@@ -1,7 +1,6 @@
 import subprocess
 subprocess.run('pip install flash-attn==2.7.0.post2 --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-import spaces
 import os
 import re
 import logging
@@ -12,14 +11,17 @@ import torch
 import gradio as gr
 from transformers import AutoModelForCausalLM, TextIteratorStreamer
 
+# 모델 및 토크나이저 로딩
 model_name = 'AIDC-AI/Ovis2-8B'
 use_thread = False
 
-# load model
-model = AutoModelForCausalLM.from_pretrained(model_name,
-                                             torch_dtype=torch.bfloat16,
-                                             multimodal_max_length=8192,
-                                             trust_remote_code=True).to(device='cuda')
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.bfloat16,
+    multimodal_max_length=8192,
+    trust_remote_code=True
+).to(device='cuda')
+
 text_tokenizer = model.get_text_tokenizer()
 visual_tokenizer = model.get_visual_tokenizer()
 streamer = TextIteratorStreamer(text_tokenizer, skip_prompt=True, skip_special_tokens=True)
@@ -45,9 +47,10 @@ def initialize_gen_kwargs():
 def submit_chat(chatbot, text_input):
     response = ''
     chatbot.append((text_input, response))
-    return chatbot ,''
+    return chatbot, ''
 
-@spaces.GPU
+@gradio.routes.no_temp_folder()
+@gradio.gpu()
 def ovis_chat(chatbot: List[List[str]], image_input: Any):
     conversations, model_inputs = prepare_inputs(chatbot, image_input)
     gen_kwargs = initialize_gen_kwargs()
@@ -72,14 +75,8 @@ def ovis_chat(chatbot: List[List[str]], image_input: Any):
 
     log_conversation(chatbot)
 
-    
 def prepare_inputs(chatbot: List[List[str]], image_input: Any):
-    # conversations = [{
-    #     "from": "system",
-    #     "value": "You are a helpful assistant, and your task is to provide reliable and structured responses to users."
-    # }]
-    conversations= []
-
+    conversations = []
     for query, response in chatbot[:-1]:
         conversations.extend([
             {"from": "human", "value": query},
@@ -90,6 +87,7 @@ def prepare_inputs(chatbot: List[List[str]], image_input: Any):
     conversations.append({"from": "human", "value": last_query})
 
     if image_input is not None:
+        # 이미지가 포함되면 첫 번째 human 메시지에 이미지 태그 추가
         for conv in conversations:
             if conv["from"] == "human":
                 conv["value"] = f'{image_placeholder}\n{conv["value"]}'
@@ -116,6 +114,7 @@ def log_conversation(chatbot):
 def clear_chat():
     return [], None, ""
 
+# 로고 SVG 로드 및 스타일 수정
 with open(f"{cur_dir}/resource/logo.svg", "r", encoding="utf-8") as svg_file:
     svg_content = svg_file.read()
 font_size = "2.5em"
@@ -125,7 +124,14 @@ html = f"""
     <span style="display: inline-block; vertical-align: middle;">{svg_content}</span>
     <span style="display: inline-block; vertical-align: middle;">{model_name.split('/')[-1]}</span>
 </p>
-<center><font size=3><b>Ovis</b> has been open-sourced on <a href='https://huggingface.co/{model_name}'>😊 Huggingface</a> and <a href='https://github.com/AIDC-AI/Ovis'>🌟 GitHub</a>. If you find Ovis useful, a like❤️ or a star🌟 would be appreciated.</font></center>
+<center>
+    <font size=3>
+        <b>Ovis</b> has been open-sourced on 
+        <a href='https://huggingface.co/{model_name}'>😊 Huggingface</a> and 
+        <a href='https://github.com/AIDC-AI/Ovis'>🌟 GitHub</a>. 
+        If you find Ovis useful, a like❤️ or a star🌟 would be appreciated.
+    </font>
+</center>
 """
 
 latex_delimiters_set = [{
@@ -158,12 +164,68 @@ latex_delimiters_set = [{
         "display": True
     }]
 
-text_input = gr.Textbox(label="prompt", placeholder="Enter your text here...", lines=1, container=False)
-with gr.Blocks(title=model_name.split('/')[-1], theme=gr.themes.Ocean()) as demo:
+text_input = gr.Textbox(label="Prompt", placeholder="Enter your text here...", lines=1, container=False)
+
+# 커스텀 CSS (배경 그라데이션, 반투명 컨테이너, 버튼 애니메이션 등)
+custom_css = """
+body {
+    background: linear-gradient(135deg, #667eea, #764ba2);
+    font-family: 'Helvetica Neue', Helvetica, Arial, sans-serif;
+    color: #333;
+    margin: 0;
+    padding: 0;
+}
+.gradio-container {
+    background: rgba(255, 255, 255, 0.95);
+    border-radius: 15px;
+    padding: 30px 40px;
+    box-shadow: 0 8px 30px rgba(0, 0, 0, 0.3);
+    margin: 40px auto;
+    max-width: 1200px;
+}
+.gradio-container h1 {
+    color: #333;
+    text-shadow: 1px 1px 2px rgba(0, 0, 0, 0.2);
+}
+.fillable { 
+    width: 95% !important; 
+    max-width: unset !important;
+}
+#examples_container {
+    margin: auto;
+    width: 90%;
+}
+#examples_row {
+    justify-content: center;
+}
+.sidebar {
+    background: rgba(255, 255, 255, 0.98);
+    border-radius: 10px;
+    padding: 20px;
+    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.2);
+}
+button, .btn {
+    background: linear-gradient(90deg, #ff8a00, #e52e71);
+    border: none;
+    color: #fff;
+    padding: 12px 24px;
+    text-transform: uppercase;
+    font-weight: bold;
+    letter-spacing: 1px;
+    border-radius: 5px;
+    cursor: pointer;
+    transition: transform 0.2s ease-in-out;
+}
+button:hover, .btn:hover {
+    transform: scale(1.05);
+}
+"""
+
+with gr.Blocks(css=custom_css, title=model_name.split('/')[-1]) as demo:
     gr.HTML(html)
     with gr.Row():
         with gr.Column(scale=3):
-            image_input = gr.Image(label="image", height=350, type="pil")
+            image_input = gr.Image(label="Image", height=350, type="pil")
             gr.Examples(
                 examples=[
                     [f"{cur_dir}/examples/ovis2_math0.jpg", "Each face of the polyhedron shown is either a triangle or a square. Each square borders 4 triangles, and each triangle borders 3 squares. The polyhedron has 6 squares. How many triangles does it have?\n\nProvide a step-by-step solution to the problem, and conclude with 'the answer is' followed by the final solution."],
@@ -178,11 +240,27 @@ with gr.Blocks(title=model_name.split('/')[-1], theme=gr.themes.Ocean()) as demo
             chatbot = gr.Chatbot(label="Ovis", layout="panel", height=600, show_copy_button=True, latex_delimiters=latex_delimiters_set)
             text_input.render()
             with gr.Row():
-                send_btn = gr.Button("Send", variant="primary")
-                clear_btn = gr.Button("Clear", variant="secondary")
+                send_btn = gr.Button("Send")
+                clear_btn = gr.Button("Clear")
 
-    send_click_event = send_btn.click(submit_chat, [chatbot, text_input], [chatbot, text_input]).then(ovis_chat,[chatbot, image_input],chatbot)
-    submit_event = text_input.submit(submit_chat, [chatbot, text_input], [chatbot, text_input]).then(ovis_chat,[chatbot, image_input],chatbot)
+    send_click_event = send_btn.click(
+        submit_chat, 
+        inputs=[chatbot, text_input], 
+        outputs=[chatbot, text_input]
+    ).then(
+        ovis_chat, 
+        inputs=[chatbot, image_input], 
+        outputs=chatbot
+    )
+    submit_event = text_input.submit(
+        submit_chat, 
+        inputs=[chatbot, text_input], 
+        outputs=[chatbot, text_input]
+    ).then(
+        ovis_chat, 
+        inputs=[chatbot, image_input], 
+        outputs=chatbot
+    )
     clear_btn.click(clear_chat, outputs=[chatbot, image_input, text_input])
 
 demo.launch()