Spaces:
Running on Zero
Running on Zero
File size: 2,999 Bytes
91ab284 c2cee1f 91ab284 c2cee1f 91ab284 c2cee1f 91ab284 c2cee1f 91ab284 c2cee1f 91ab284 c2cee1f 91ab284 c2cee1f 91ab284 c2cee1f 91ab284 c2cee1f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | import spaces
import torch
from transformers import AutoModel, AutoProcessor, PreTrainedModel
# ============================================================
# PATCH KOMPATIBILITAS
# transformers >= 5.9 mengubah _keys_to_ignore_on_load_unexpected
# dari list-append jadi set-union. Remote code ZDTaichu5.0 masih
# pakai list, sehingga muncul: TypeError: unsupported operand
# type(s) for |: 'list' and 'set'
# Patch ini paksa konversi ke set sebelum method asli dijalankan.
# ============================================================
_orig_adjust = PreTrainedModel._adjust_missing_and_unexpected_keys
def _patched_adjust(self, loading_info):
attr = getattr(self, "_keys_to_ignore_on_load_unexpected", None)
if isinstance(attr, (list, tuple)):
self._keys_to_ignore_on_load_unexpected = set(attr)
attr2 = getattr(self, "_keys_to_ignore_on_load_missing", None)
if isinstance(attr2, (list, tuple)):
self._keys_to_ignore_on_load_missing = set(attr2)
return _orig_adjust(self, loading_info)
PreTrainedModel._adjust_missing_and_unexpected_keys = _patched_adjust
# ============================================================
# LOAD MODEL DI MODULE LEVEL
# ZeroGPU butuh load di luar @spaces.GPU.
# device_map="cuda" (bukan "auto") supaya ZeroGPU bisa emulate.
# ============================================================
MODEL_ID = "TaichuAI/ZDTaichu5.0-9B"
processor = AutoProcessor.from_pretrained(
MODEL_ID,
trust_remote_code=True,
)
model = AutoModel.from_pretrained(
MODEL_ID,
torch_dtype=torch.bfloat16,
device_map="cuda",
trust_remote_code=True,
attn_implementation="sdpa",
).eval()
# ============================================================
# INFERENCE FUNCTION
# ============================================================
@spaces.GPU(duration=60)
def generate(image, text):
if image is None:
content = [{"type": "text", "text": text}]
else:
content = [
{"type": "image", "image": image},
{"type": "text", "text": text},
]
messages = [{"role": "user", "content": content}]
inputs = processor.from_messages(messages, return_tensors="pt").to("cuda")
with torch.inference_mode():
output_ids = model.generate(
**inputs,
max_new_tokens=512,
do_sample=False,
)
generated_ids = output_ids[:, inputs["input_ids"].shape[1]:]
return processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
# ============================================================
# GRADIO UI
# ============================================================
import gradio as gr
demo = gr.Interface(
fn=generate,
inputs=[
gr.Image(type="pil", label="Gambar (opsional)"),
gr.Textbox(label="Pertanyaan", lines=3),
],
outputs=gr.Textbox(label="Jawaban", lines=10),
title="ZDTaichu5.0-9B Demo",
description="Multimodal VLM — teks + gambar. Jalan di ZeroGPU.",
)
demo.launch() |