app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import json
import torch
import requests
import time
import random
from PIL import Image
from typing import Union
import os
import base64
from together import Together
import pathlib

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device}" if device != "cpu" else "Using CPU")

def _load_model():
    tokenizer = AutoTokenizer.from_pretrained("vikhyatk/moondream2", trust_remote_code=True, revision="2024-05-08", torch_dtype=(torch.bfloat16 if device == 'cuda' else torch.float32))
    model = AutoModelForCausalLM.from_pretrained("vikhyatk/moondream2", device_map=device, trust_remote_code=True, revision="2024-05-08")
    return (model, tokenizer)

class MoonDream():
    def __init__(self, model=None, tokenizer=None):
        self.model, self.tokenizer = (model, tokenizer)
        if not model or model is None or not tokenizer or tokenizer is None:
            self.model, self.tokenizer = _load_model()
        self.device = device
        self.model.to(self.device)
    def __call__(self, question, imgs):
        imn = 0
        for img in imgs:
            img = self.model.encode_image(img)
            res = self.model.answer_question(question=question, image_embeds=img, tokenizer=self.tokenizer)
            yield res
        return

md = MoonDream()

SYSTEM_PROMPT = "You are Llama 3 70b. You have been given access to Moondream 2 for VQA when given images. When you have a question about an image, simple start your response with the text, '@question\\nMy question?'. When you do this, the request will be sent to Moondream 2. User can see this happening if they turn debug on, so be professional and stay on topic. Any chat from anyone starting with @answer is the answer to last question asked. If something appears out of sync, ask User to clear the chat."

def _respond_one(question, img):
  txt = ""
  yield (txt := txt + MoonDream()(question, [img]))
  return txt

def respond_batch(question, **imgs):
  md = MoonDream()
  for img in imgs.values():
    res = md(question, img)
    for r in res:
      yield r
    yield "\n\n\n\n\n\n"
  return

def dual_images(img1: Image):
  # Ran once for each img to it's respective output. Output should be detailed str of description/feature extraction/interrogation.
  md = MoonDream()
  res = md("Describe the image in plain english ", [img1])
  txt = ""
  for r in res:
    yield (txt := txt + r)
  return

import os

def merge_descriptions_to_prompt(mi, d1, d2):
  from together import Together
  tog = Together(api_key=os.getenv("TOGETHER_KEY"))
  res = tog.completions.create(prompt=f"""Describe what would result if the following two descriptions were describing one thing.
### Description 1:
```text
{d1}
```
### Description 2:
```text
{d2}
```
Merge-Specific Instructions:
```text
{mi}
```
Ensure you end your output with ```\\n
---
Complete Description:
```text""", model="meta-llama/Meta-Llama-3-70B", stop=["```"], max_tokens=1024)
  return res.choices[0].text.split("```")[0]

def xform_image_description(img, inst):
  #md = MoonDream()
  from together import Together
  desc = dual_images(img)
  tog = Together(api_key=os.getenv("TOGETHER_KEY"))
  prompt=f"""Describe the image in aggressively verbose detail. I must know every freckle upon a man's brow and each blade of the grass intimately.\nDescription: ```text\n{desc}\n```\nInstructions:\n```text\n{inst}\n```\n\n\n---\nDetailed Description:\n```text"""
  res = tog.completions.create(prompt=prompt, model="meta-llama/Meta-Llama-3-70B", stop=["```"], max_tokens=1024)
  return res.choices[0].text[len(prompt):].split("```")[0]

def simple_desc(img, prompt):
  import base64
  gen = md(prompt, [img])
  total = ""
  for resp in gen:
    print(total := total + resp)
  img.resize((192,192)).save("tmp.png")
  bts = False
  with open("tmp.png", "rb") as f:
    bts = f.read()
  if bts:
    os.remove("tmp.png")
  res = {
    'image_b64': base64.b64encode(bts).decode('utf-8'),
    'description': total,
  }
  return total, res

ifc_imgprompt2text = gr.Interface(simple_desc, inputs=[gr.Image(label="input", type="pil"), gr.Textbox(label="prompt")], outputs=[gr.Textbox(label="description"), gr.JSON(label="json")])

def chat(inpt, mess):
  from together import Together
  print(inpt, mess)
  if mess is None:
    mess = []

  tog = Together(api_key=os.getenv("TOGETHER_KEY"))
  messages = [
    {
      'role': 'system',
      'content': SYSTEM_PROMPT
    },
    {
      'role': 'user',
      'content': inpt
    }
  ]
  for cht in mess:
    print(cht)
  res = tog.chat.completions.create(
  messages=messages,
  model="meta-llama/Llama-3-70b-chat-hf", stop=["<|eot_id|>"], stream=True)
  txt = ""
  for pk in res:
    print(pk)
    txt += pk.choices[0].delta.content
    #mess[-1][-2] += pk.choices[0].delta.content
    yield txt #, json.dumps(messages)#mess#, json.dumps(messages)

chatbot = gr.Chatbot(
  [],
  elem_id="chatbot",
  bubble_full_width=False,
  sanitize_html=False,
  show_copy_button=True,
  avatar_images=[
    pathlib.Path("image.jpeg"), 
    pathlib.Path("image2.jpeg")
])

with gr.TabbedInterface([ifc_imgprompt2text, gr.ChatInterface(chat, chatbot=chatbot, submit_btn=gr.Button(scale=1))], ["Prompt & Image 2 Text", "Chat w/ Llama 3 70b & Moondream 2"]) as ifc:
  ifc.launch(share=False, debug=True)