Spaces:
Running
Running
#!/usr/bin/env python | |
import os | |
import re | |
import tempfile | |
from collections.abc import Iterator | |
from threading import Thread | |
import cv2 | |
import gradio as gr | |
import spaces | |
import torch | |
from loguru import logger | |
from PIL import Image | |
from transformers import AutoProcessor, Gemma3ForConditionalGeneration, TextIteratorStreamer | |
# [CSV/TXT 분석용] | |
import pandas as pd | |
################################################## | |
# 전체 전문을 넘기되, 너무 클 경우 잘라내기 위한 상수 | |
################################################## | |
MAX_CONTENT_CHARS = 8000 # 예: 8000자 초과 시 잘라냄 | |
model_id = os.getenv("MODEL_ID", "google/gemma-3-27b-it") | |
processor = AutoProcessor.from_pretrained(model_id, padding_side="left") | |
model = Gemma3ForConditionalGeneration.from_pretrained( | |
model_id, device_map="auto", torch_dtype=torch.bfloat16, attn_implementation="eager" | |
) | |
MAX_NUM_IMAGES = int(os.getenv("MAX_NUM_IMAGES", "5")) | |
################################################## | |
# CSV/TXT 전문 처리 함수 | |
################################################## | |
def analyze_csv_file(path: str) -> str: | |
""" | |
CSV 파일 전체를 문자열로 변환하여 리턴. | |
너무 길면 MAX_CONTENT_CHARS까지만 잘라냄. | |
""" | |
try: | |
df = pd.read_csv(path) | |
df_str = df.to_string() | |
if len(df_str) > MAX_CONTENT_CHARS: | |
df_str = df_str[:MAX_CONTENT_CHARS] + "\n...(truncated)..." | |
return ( | |
f"**[CSV File: {os.path.basename(path)}]**\n\n" | |
f"{df_str}" | |
) | |
except Exception as e: | |
return f"Failed to read CSV ({os.path.basename(path)}): {str(e)}" | |
def analyze_txt_file(path: str) -> str: | |
""" | |
TXT 파일 전체 내용을 읽어서 모델에 넘김. | |
너무 길면 MAX_CONTENT_CHARS까지만 잘라냄. | |
""" | |
try: | |
with open(path, "r", encoding="utf-8") as f: | |
text = f.read() | |
if len(text) > MAX_CONTENT_CHARS: | |
text = text[:MAX_CONTENT_CHARS] + "\n...(truncated)..." | |
return ( | |
f"**[TXT File: {os.path.basename(path)}]**\n\n" | |
f"{text}" | |
) | |
except Exception as e: | |
return f"Failed to read TXT ({os.path.basename(path)}): {str(e)}" | |
################################################## | |
# 기존 미디어 파일 검사 로직 (이미지/비디오) | |
################################################## | |
def count_files_in_new_message(paths: list[str]) -> tuple[int, int]: | |
image_count = 0 | |
video_count = 0 | |
for path in paths: | |
if path.endswith(".mp4"): | |
video_count += 1 | |
else: | |
image_count += 1 | |
return image_count, video_count | |
def count_files_in_history(history: list[dict]) -> tuple[int, int]: | |
image_count = 0 | |
video_count = 0 | |
for item in history: | |
if item["role"] != "user" or isinstance(item["content"], str): | |
continue | |
if item["content"][0].endswith(".mp4"): | |
video_count += 1 | |
else: | |
image_count += 1 | |
return image_count, video_count | |
def validate_media_constraints(message: dict, history: list[dict]) -> bool: | |
""" | |
- 비디오 1개 초과 불가 | |
- 비디오/이미지 혼합 불가 | |
- 이미지 개수 MAX_NUM_IMAGES 초과 불가 | |
- <image> 태그가 있으면 태그 수와 이미지 수 일치 | |
CSV, TXT, PDF 등은 여기서 제한하지 않음. | |
""" | |
media_files = [] | |
for f in message["files"]: | |
# mp4나 대표 이미지 확장자만 검사 | |
# (파일명에 .png / .jpg / .gif / .webp 등 있을 때) | |
if f.endswith(".mp4") or re.search(r"\.(png|jpg|jpeg|gif|webp)$", f, re.IGNORECASE): | |
media_files.append(f) | |
new_image_count, new_video_count = count_files_in_new_message(media_files) | |
history_image_count, history_video_count = count_files_in_history(history) | |
image_count = history_image_count + new_image_count | |
video_count = history_video_count + new_video_count | |
if video_count > 1: | |
gr.Warning("Only one video is supported.") | |
return False | |
if video_count == 1: | |
if image_count > 0: | |
gr.Warning("Mixing images and videos is not allowed.") | |
return False | |
if "<image>" in message["text"]: | |
gr.Warning("Using <image> tags with video files is not supported.") | |
return False | |
if video_count == 0 and image_count > MAX_NUM_IMAGES: | |
gr.Warning(f"You can upload up to {MAX_NUM_IMAGES} images.") | |
return False | |
if "<image>" in message["text"] and message["text"].count("<image>") != new_image_count: | |
gr.Warning("The number of <image> tags in the text does not match the number of images.") | |
return False | |
return True | |
################################################## | |
# 비디오 처리 | |
################################################## | |
def downsample_video(video_path: str) -> list[tuple[Image.Image, float]]: | |
vidcap = cv2.VideoCapture(video_path) | |
fps = vidcap.get(cv2.CAP_PROP_FPS) | |
total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
frame_interval = int(fps / 3) | |
frames = [] | |
for i in range(0, total_frames, frame_interval): | |
vidcap.set(cv2.CAP_PROP_POS_FRAMES, i) | |
success, image = vidcap.read() | |
if success: | |
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
pil_image = Image.fromarray(image) | |
timestamp = round(i / fps, 2) | |
frames.append((pil_image, timestamp)) | |
vidcap.release() | |
return frames | |
def process_video(video_path: str) -> list[dict]: | |
content = [] | |
frames = downsample_video(video_path) | |
for frame in frames: | |
pil_image, timestamp = frame | |
with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file: | |
pil_image.save(temp_file.name) | |
content.append({"type": "text", "text": f"Frame {timestamp}:"}) | |
content.append({"type": "image", "url": temp_file.name}) | |
logger.debug(f"{content=}") | |
return content | |
################################################## | |
# interleaved <image> 태그 처리 | |
################################################## | |
def process_interleaved_images(message: dict) -> list[dict]: | |
logger.debug(f"{message['files']=}") | |
parts = re.split(r"(<image>)", message["text"]) | |
logger.debug(f"{parts=}") | |
content = [] | |
image_index = 0 | |
for part in parts: | |
logger.debug(f"{part=}") | |
if part == "<image>": | |
content.append({"type": "image", "url": message["files"][image_index]}) | |
logger.debug(f"file: {message['files'][image_index]}") | |
image_index += 1 | |
elif part.strip(): | |
content.append({"type": "text", "text": part.strip()}) | |
elif isinstance(part, str) and part != "<image>": | |
content.append({"type": "text", "text": part}) | |
logger.debug(f"{content=}") | |
return content | |
################################################## | |
# CSV, TXT 파일도 전문을 LLM에 넘기도록 | |
################################################## | |
def process_new_user_message(message: dict) -> list[dict]: | |
""" | |
- mp4 -> 비디오 처리 | |
- 이미지 -> interleaved or multiple | |
- CSV -> 전체 df.to_string() (너무 길면 잘라냄) | |
- TXT -> 전체 text (너무 길면 잘라냄) | |
""" | |
if not message["files"]: | |
return [{"type": "text", "text": message["text"]}] | |
# 확장자별 분류 | |
video_files = [f for f in message["files"] if f.endswith(".mp4")] | |
image_files = [f for f in message["files"] if re.search(r"\.(png|jpg|jpeg|gif|webp)$", f, re.IGNORECASE)] | |
csv_files = [f for f in message["files"] if f.lower().endswith(".csv")] | |
txt_files = [f for f in message["files"] if f.lower().endswith(".txt")] | |
# 사용자 텍스트 | |
content_list = [{"type": "text", "text": message["text"]}] | |
# CSV 전문 | |
for csv_path in csv_files: | |
csv_analysis = analyze_csv_file(csv_path) | |
content_list.append({"type": "text", "text": csv_analysis}) | |
# TXT 전문 | |
for txt_path in txt_files: | |
txt_analysis = analyze_txt_file(txt_path) | |
content_list.append({"type": "text", "text": txt_analysis}) | |
# 비디오 | |
if video_files: | |
content_list += process_video(video_files[0]) | |
return content_list | |
# interleaved 이미지 | |
if "<image>" in message["text"]: | |
return process_interleaved_images(message) | |
# 일반 이미지(여러 장) | |
if image_files: | |
for img_path in image_files: | |
content_list.append({"type": "image", "url": img_path}) | |
return content_list | |
################################################## | |
# history -> LLM 메시지 변환 | |
################################################## | |
def process_history(history: list[dict]) -> list[dict]: | |
messages = [] | |
current_user_content: list[dict] = [] | |
for item in history: | |
if item["role"] == "assistant": | |
if current_user_content: | |
messages.append({"role": "user", "content": current_user_content}) | |
current_user_content = [] | |
messages.append({"role": "assistant", "content": [{"type": "text", "text": item["content"]}]}) | |
else: | |
content = item["content"] | |
if isinstance(content, str): | |
current_user_content.append({"type": "text", "text": content}) | |
else: | |
current_user_content.append({"type": "image", "url": content[0]}) | |
return messages | |
################################################## | |
# 메인 추론 함수 | |
################################################## | |
def run(message: dict, history: list[dict], system_prompt: str = "", max_new_tokens: int = 512) -> Iterator[str]: | |
if not validate_media_constraints(message, history): | |
yield "" | |
return | |
messages = [] | |
if system_prompt: | |
messages.append({"role": "system", "content": [{"type": "text", "text": system_prompt}]}) | |
messages.extend(process_history(history)) | |
messages.append({"role": "user", "content": process_new_user_message(message)}) | |
inputs = processor.apply_chat_template( | |
messages, | |
add_generation_prompt=True, | |
tokenize=True, | |
return_dict=True, | |
return_tensors="pt", | |
).to(device=model.device, dtype=torch.bfloat16) | |
streamer = TextIteratorStreamer(processor, timeout=30.0, skip_prompt=True, skip_special_tokens=True) | |
generate_kwargs = dict( | |
inputs, | |
streamer=streamer, | |
max_new_tokens=max_new_tokens, | |
) | |
t = Thread(target=model.generate, kwargs=generate_kwargs) | |
t.start() | |
output = "" | |
for delta in streamer: | |
output += delta | |
yield output | |
################################################## | |
# 예시 목록 (기존) | |
################################################## | |
examples = [ | |
[ | |
{ | |
"text": "I need to be in Japan for 10 days, going to Tokyo, Kyoto and Osaka. Think about number of attractions in each of them and allocate number of days to each city. Make public transport recommendations.", | |
"files": [], | |
} | |
], | |
[ | |
{ | |
"text": "Write the matplotlib code to generate the same bar chart.", | |
"files": ["assets/additional-examples/barchart.png"], | |
} | |
], | |
[ | |
{ | |
"text": "What is odd about this video?", | |
"files": ["assets/additional-examples/tmp.mp4"], | |
} | |
], | |
[ | |
{ | |
"text": "I already have this supplement <image> and I want to buy this one <image>. Any warnings I should know about?", | |
"files": ["assets/additional-examples/pill1.png", "assets/additional-examples/pill2.png"], | |
} | |
], | |
[ | |
{ | |
"text": "Write a poem inspired by the visual elements of the images.", | |
"files": ["assets/sample-images/06-1.png", "assets/sample-images/06-2.png"], | |
} | |
], | |
[ | |
{ | |
"text": "Compose a short musical piece inspired by the visual elements of the images.", | |
"files": [ | |
"assets/sample-images/07-1.png", | |
"assets/sample-images/07-2.png", | |
"assets/sample-images/07-3.png", | |
"assets/sample-images/07-4.png", | |
], | |
} | |
], | |
[ | |
{ | |
"text": "Write a short story about what might have happened in this house.", | |
"files": ["assets/sample-images/08.png"], | |
} | |
], | |
[ | |
{ | |
"text": "Create a short story based on the sequence of images.", | |
"files": [ | |
"assets/sample-images/09-1.png", | |
"assets/sample-images/09-2.png", | |
"assets/sample-images/09-3.png", | |
"assets/sample-images/09-4.png", | |
"assets/sample-images/09-5.png", | |
], | |
} | |
], | |
[ | |
{ | |
"text": "Describe the creatures that would live in this world.", | |
"files": ["assets/sample-images/10.png"], | |
} | |
], | |
[ | |
{ | |
"text": "Read text in the image.", | |
"files": ["assets/additional-examples/1.png"], | |
} | |
], | |
[ | |
{ | |
"text": "When is this ticket dated and how much did it cost?", | |
"files": ["assets/additional-examples/2.png"], | |
} | |
], | |
[ | |
{ | |
"text": "Read the text in the image into markdown.", | |
"files": ["assets/additional-examples/3.png"], | |
} | |
], | |
[ | |
{ | |
"text": "Evaluate this integral.", | |
"files": ["assets/additional-examples/4.png"], | |
} | |
], | |
[ | |
{ | |
"text": "caption this image", | |
"files": ["assets/sample-images/01.png"], | |
} | |
], | |
[ | |
{ | |
"text": "What's the sign says?", | |
"files": ["assets/sample-images/02.png"], | |
} | |
], | |
[ | |
{ | |
"text": "Compare and contrast the two images.", | |
"files": ["assets/sample-images/03.png"], | |
} | |
], | |
[ | |
{ | |
"text": "List all the objects in the image and their colors.", | |
"files": ["assets/sample-images/04.png"], | |
} | |
], | |
[ | |
{ | |
"text": "Describe the atmosphere of the scene.", | |
"files": ["assets/sample-images/05.png"], | |
} | |
], | |
] | |
################################################## | |
# Gradio ChatInterface | |
################################################## | |
demo = gr.ChatInterface( | |
fn=run, | |
type="messages", | |
chatbot=gr.Chatbot(type="messages", scale=1, allow_tags=["image"]), | |
# 여기서 WEBP를 포함한 모든 이미지, mp4, csv, txt, pdf 허용 | |
textbox=gr.MultimodalTextbox( | |
file_types=["image/*", ".mp4", ".csv", ".txt", ".pdf"], | |
file_count="multiple", | |
autofocus=True | |
), | |
multimodal=True, | |
additional_inputs=[ | |
gr.Textbox( | |
label="System Prompt", | |
value=( | |
"You are a deeply thoughtful AI. Consider problems thoroughly and derive " | |
"correct solutions through systematic reasoning. Please answer in korean." | |
) | |
), | |
gr.Slider( | |
label="Max New Tokens", | |
minimum=100, | |
maximum=8000, | |
step=50, | |
value=2000 | |
), | |
], | |
stop_btn=False, | |
title="Gemma 3 27B IT", | |
examples=examples, | |
run_examples_on_click=False, | |
cache_examples=False, | |
css_paths="style.css", | |
delete_cache=(1800, 1800), | |
) | |
if __name__ == "__main__": | |
demo.launch() | |