Update app.py

app.py

@@ -33,56 +33,6 @@ model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True
     )
 
-# Chatbot tab function
-@spaces.GPU()
-def stream_chat(
-    message: str,
-    history: list,
-    system_prompt: str,
-    temperature: float = 0.8,
-    max_new_tokens: int = 1024,
-    top_p: float = 1.0,
-    top_k: int = 20,
-    penalty: float = 1.2,
-):
-    print(f'message: {message}')
-    print(f'history: {history}')
-
-    conversation = [
-        {"role": "system", "content": system_prompt}
-    ]
-    for prompt, answer in history:
-        conversation.extend([
-            {"role": "user", "content": prompt},
-            {"role": "assistant", "content": answer},
-        ])
-
-    conversation.append({"role": "user", "content": message})
-
-    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(model.device)
-
-    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
-
-    generate_kwargs = dict(
-        input_ids=input_ids,
-        max_new_tokens = max_new_tokens,
-        do_sample = False if temperature == 0 else True,
-        top_p = top_p,
-        top_k = top_k,
-        temperature = temperature,
-        eos_token_id=[128001,128008,128009],
-        streamer=streamer,
-    )
-
-    with torch.no_grad():
-        thread = Thread(target=model.generate, kwargs=generate_kwargs)
-        thread.start()
-
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        yield buffer
-
 # Vision model setup
 models = {
     "justinj92/phi-35-vision-burberry": AutoModelForCausalLM.from_pretrained("justinj92/phi-35-vision-burberry", trust_remote_code=True, torch_dtype="auto", _attn_implementation="flash_attention_2").cuda().eval()
@@ -184,20 +134,20 @@ footer = """
 """
 
 # Gradio app with two tabs
-with gr.Blocks(css=CSS, theme="small_and_pretty") as demo:
+with gr.Blocks(css=CSS, theme=gr.themes.Default(primary_hue=gr.themes.colors.red, secondary_hue=gr.themes.colors.pink)) as demo:
     gr.HTML(TITLE)
     gr.HTML(EXPLANATION)
     with gr.Tab("Burberry Vision"):
         with gr.Row():
-            input_img = gr.Image(label="Input Picture")
+            input_img = gr.Image(label="Upload a Burberry Product Image")
         with gr.Row():
             model_selector = gr.Dropdown(choices=list(models.keys()), label="Model", value="justinj92/phi-35-vision-burberry")
         # with gr.Row():
         #     text_input = gr.Textbox(label="Question")
         with gr.Row():
-            submit_btn = gr.Button(value="Submit")
+            submit_btn = gr.Button(value="Tell me about this product")
         with gr.Row():
-            output_text = gr.Textbox(label="Output Text")
+            output_text = gr.Textbox(label="Product Info")
 
         submit_btn.click(stream_vision, [input_img, model_selector], [output_text])