Spaces:
Paused
Paused
| import os | |
| import threading | |
| import time | |
| import gradio as gr | |
| import numpy as np | |
| import spaces | |
| import torch | |
| from huggingface_hub import hf_hub_download, snapshot_download | |
| from omegaconf import OmegaConf | |
| from PIL import Image | |
| from transformers import StoppingCriteriaList | |
| from minigpt4.conversation.conversation import ( | |
| CONV_VISION_Vicuna0, Chat, StoppingCriteriaSub) | |
| from minigpt4.models.minigpt4 import MiniGPT4 | |
| from minigpt4.processors.blip_processors import Blip2ImageEvalProcessor | |
| # Vicuna v0 normally needs LLaMA-1 base weights plus a delta, and LLaMA-1 was never | |
| # publicly distributed. These are the authors' own prepared weights, which is the | |
| # only reason this model can still be assembled at all. | |
| VICUNA_REPO = "Vision-CAIR/vicuna-7b" | |
| # The trained linear projection. Everything else (ViT-G, Q-Former) stays frozen. | |
| CKPT_REPO, CKPT_FILE = "ckpt/minigpt4-7B", "prerained_minigpt4_7b.pth" | |
| MODEL = None | |
| CHAT = None | |
| VISION_OK = None | |
| LOCK = threading.Lock() | |
| _ON_GPU = False | |
| def build(): | |
| global MODEL, CHAT, VISION_OK | |
| vicuna = snapshot_download(VICUNA_REPO) | |
| ckpt = hf_hub_download(CKPT_REPO, CKPT_FILE) | |
| cfg = OmegaConf.load("minigpt4/configs/models/minigpt4_vicuna0.yaml").model | |
| cfg.llama_model = vicuna | |
| cfg.ckpt = ckpt | |
| cfg.low_resource = False # 24 GB of A10G, no need for the 8-bit path | |
| cfg.prompt_template = "###Human: {} ###Assistant: " | |
| cfg.max_txt_len = 160 | |
| cfg.end_sym = "###" | |
| MODEL = MiniGPT4.from_config(cfg).eval() | |
| vis = Blip2ImageEvalProcessor.from_config( | |
| OmegaConf.create({"name": "blip2_image_eval", "image_size": 224})) | |
| # "###" and "##" as stop sequences, straight from the reference demo. | |
| stops = [torch.tensor(ids) for ids in ([835], [2277, 29937])] | |
| CHAT = Chat(MODEL, vis, device="cpu", | |
| stopping_criteria=StoppingCriteriaList([StoppingCriteriaSub(stops=stops)])) | |
| print("[mg4] model built on CPU", flush=True) | |
| def to_gpu(): | |
| """ZeroGPU only exposes CUDA inside a @spaces.GPU call, so the move happens on | |
| first use rather than at import.""" | |
| global _ON_GPU | |
| if _ON_GPU: | |
| return | |
| dev = torch.device("cuda") | |
| MODEL.to(dev) | |
| CHAT.device = dev | |
| CHAT.stopping_criteria = StoppingCriteriaList( | |
| [StoppingCriteriaSub(stops=[torch.tensor(i).to(dev) for i in ([835], [2277, 29937])])]) | |
| _ON_GPU = True | |
| print("[mg4] moved to CUDA", flush=True) | |
| def _embed(img): | |
| lst = [] | |
| conv = CONV_VISION_Vicuna0.copy() | |
| CHAT.upload_img(img, conv, lst) | |
| CHAT.encode_img(lst) | |
| return lst[0].detach().float().cpu().numpy() | |
| def selftest(): | |
| """The CPU ggml port returned bit-identical embeddings for any two images, which | |
| is what made every answer confabulated. Prove this build does not, and time the | |
| encode while we are here so the header can quote a real number.""" | |
| to_gpu() | |
| black = Image.new("RGB", (448, 448), (0, 0, 0)) | |
| patt = Image.fromarray( | |
| (np.indices((448, 448)).sum(0) % 256).astype("uint8")).convert("RGB") | |
| a = _embed(black) | |
| t = time.time() | |
| b = _embed(patt) | |
| enc = time.time() - t | |
| diff = float(np.abs(a - b).max()) | |
| ok = diff > 1e-4 | |
| print(f"[mg4] SELFTEST vision responds to input: {ok} (maxdiff={diff:.6f} " | |
| f"shape={a.shape} encode={enc:.2f}s)", flush=True) | |
| return ok, enc | |
| def example_run(image, prompt): | |
| """The example row carries only the two inputs, so the sliders keep their | |
| defaults. Routed through chat() so a cached example is the same code path | |
| the Send button runs.""" | |
| return chat(image, prompt, [], 300, 0.7, 0.9) | |
| def chat(image, prompt, history, max_new, temperature, top_p): | |
| history = history or [] | |
| if image is None: | |
| raise gr.Error("Upload an image first. This model always answers about an image.") | |
| if not (prompt or "").strip(): | |
| raise gr.Error("Type a question about the image.") | |
| if not LOCK.acquire(blocking=False): | |
| raise gr.Error("Another request is running. Try again in a moment.") | |
| try: | |
| to_gpu() | |
| t0 = time.time() | |
| conv = CONV_VISION_Vicuna0.copy() | |
| img_list = [] | |
| CHAT.upload_img(image, conv, img_list) | |
| CHAT.encode_img(img_list) | |
| enc = time.time() - t0 | |
| CHAT.ask(prompt, conv) | |
| t1 = time.time() | |
| answer = CHAT.answer(conv=conv, img_list=img_list, | |
| num_beams=1, temperature=float(temperature), | |
| top_p=float(top_p), max_new_tokens=int(max_new), | |
| max_length=2000)[0] | |
| dt = time.time() - t1 | |
| answer = answer.replace("<Img>", "").replace("</Img>", "").strip() | |
| note = f"\n\n`encode {enc:.1f}s · generate {dt:.1f}s`" | |
| return history + [{"role": "user", "content": prompt}, | |
| {"role": "assistant", "content": answer + note}], "" | |
| finally: | |
| LOCK.release() | |
| CSS = """ | |
| footer{display:none} | |
| .gap{gap:0} | |
| .block{padding:2px} | |
| #img .upload-container,#img .file-preview{max-height:70px} | |
| #img img{max-height:150px;object-fit:contain} | |
| #hdr{max-height:32px} | |
| #hdr p{margin:0} | |
| :where(button,input,textarea,select,[role="button"],[tabindex]):focus-visible{ | |
| outline:2px solid #7c5cff!important;outline-offset:2px} | |
| @media (prefers-reduced-motion:reduce){*,*::before,*::after{ | |
| animation-duration:.01ms!important;transition-duration:.01ms!important}} | |
| """ | |
| build() | |
| VISION_OK, ENC_S = selftest() | |
| # One line, resolved backend and measured cost, so the wall clock is predictable | |
| # before anyone uploads anything. Generation is ~9 ms/token on this hardware. | |
| GPU_NAME = "ZeroGPU A10G" | |
| # Sub-second encodes read as "0.0s" at one decimal, which looks like broken telemetry. | |
| ENC_TXT = f"{ENC_S * 1000:.0f}ms" if ENC_S < 0.5 else f"{ENC_S:.1f}s" | |
| HEADER = (f"**MiniGPT-4 7B** · uncensored captioning · Vicuna v0 · {GPU_NAME} · " | |
| f"~{ENC_TXT} encode + ~{300 * 0.009:.0f}s per 300 tokens" | |
| + ("" if VISION_OK else " · ⚠ **vision self-test FAILED, answers are " | |
| "not grounded in the image**")) | |
| with gr.Blocks(fill_width=True) as demo: | |
| gr.Markdown(HEADER, elem_id="hdr") | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| img = gr.Image(type="pil", show_label=False, elem_id="img", height=150) | |
| box = gr.Chatbot(show_label=False, height=230) | |
| with gr.Row(): | |
| msg = gr.Textbox(placeholder="Ask about the image...", show_label=False, | |
| container=False, lines=2, max_lines=2, scale=5) | |
| go = gr.Button("Send", variant="primary", scale=1) | |
| wipe = gr.Button("New conversation", size="sm") | |
| with gr.Column(scale=1): | |
| gr.Markdown( | |
| "Vicuna v0 predates the refusal training that arrived later in 2023, so " | |
| "it describes what it is shown instead of declining. It also " | |
| "hallucinates freely. A 2023 model, not a current one." | |
| ) | |
| with gr.Accordion("Settings", open=False): | |
| max_new = gr.Slider(32, 512, value=300, step=32, | |
| label="Max new tokens (300 ≈ 3s)") | |
| temperature = gr.Slider(0.1, 1.5, value=0.7, step=0.05, | |
| label="Temperature", | |
| info="Higher wanders further from the image.") | |
| top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top P") | |
| go.click(chat, [img, msg, box, max_new, temperature, top_p], [box, msg]) | |
| msg.submit(chat, [img, msg, box, max_new, temperature, top_p], [box, msg]) | |
| wipe.click(lambda: ([], ""), outputs=[box, msg], queue=False) | |
| # Exactly one row. Every row is a full pipeline run if caching is ever enabled. | |
| # With no fn the example only fills the inputs and the visitor still has to | |
| # press Send, which reads as a broken example. Give it the handler and cache | |
| # it. Lazy rather than eager because caching at startup would need the GPU | |
| # before ZeroGPU has handed one over. | |
| gr.Examples(examples=[["examples/cats.jpg", "Describe this image in detail."]], | |
| inputs=[img, msg], outputs=[box, msg], | |
| fn=example_run, cache_examples=True, cache_mode="lazy", | |
| label="Example") | |
| demo.queue(default_concurrency_limit=1, max_size=8).launch(css=CSS, show_error=True) | |