|
import gradio as gr |
|
import spaces |
|
from PIL import Image |
|
import torch |
|
from transformers import AutoModelForCausalLM, AutoProcessor |
|
import requests |
|
import json |
|
from io import BytesIO |
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
model = AutoModelForCausalLM.from_pretrained("MiaoshouAI/Florence-2-base-PromptGen-v1.5", trust_remote_code=True).to(device) |
|
processor = AutoProcessor.from_pretrained("MiaoshouAI/Florence-2-base-PromptGen-v1.5", trust_remote_code=True) |
|
|
|
SERVER_URL = 'http://43.156.72.113:8188' |
|
FETCH_TASKS_URL = SERVER_URL + '/fetch/' |
|
UPDATE_TASK_STATUS_URL = SERVER_URL + '/update/' |
|
|
|
def fetch_task(category, fetch_all=False): |
|
params = {'fetch_all': 'true' if fetch_all else 'false'} |
|
response = requests.post(FETCH_TASKS_URL + category, params=params) |
|
if response.status_code == 200: |
|
return response.json() |
|
else: |
|
print(f"Failed to fetch tasks: {response.status_code} - {response.text}") |
|
return None |
|
|
|
def update_task_status(category, task_id, status, result=None): |
|
data = {'status': status} |
|
if result: |
|
data['result'] = result |
|
|
|
response = requests.post(UPDATE_TASK_STATUS_URL + category + f'/{task_id}', json=data) |
|
if response.status_code == 200: |
|
print(f"Task {task_id} updated successfully: {json.dumps(response.json(), indent=4)}") |
|
else: |
|
print(f"Failed to update task {task_id}: {response.status_code} - {response.text}") |
|
|
|
|
|
@spaces.GPU(duration=150) |
|
def infer(): |
|
|
|
img2text_tasks = fetch_task('img2text', fetch_all=True) |
|
|
|
if not img2text_tasks: |
|
return "No tasks found." |
|
|
|
for task in img2text_tasks: |
|
try: |
|
image_url = task['content']['url'] |
|
prompt = task['content']['prompt'] |
|
|
|
image_response = requests.get(image_url) |
|
image = Image.open(BytesIO(image_response.content)).convert("RGB") |
|
|
|
max_size = 256 |
|
width, height = image.size |
|
if width > height: |
|
new_width = max_size |
|
new_height = int((new_width / width) * height) |
|
else: |
|
new_height = max_size |
|
new_width = int((new_height / height) * width) |
|
image = image.resize((new_width, new_height), Image.LANCZOS) |
|
|
|
inputs = processor(text=prompt, images=image, return_tensors="pt").to(device) |
|
generated_ids = model.generate( |
|
input_ids=inputs["input_ids"], |
|
pixel_values=inputs["pixel_values"], |
|
max_new_tokens=1024, |
|
do_sample=False, |
|
num_beams=3 |
|
) |
|
|
|
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0] |
|
parsed_answer = processor.post_process_generation(generated_text, task=prompt, image_size=(image.width, image.height)) |
|
|
|
update_task_status('img2text', task['id'], 'Successed', {"text": parsed_answer}) |
|
except Exception as e: |
|
print(f"Error processing task {task['id']}: {e}") |
|
update_task_status('img2text', task['id'], 'Failed', {"error": str(e)}) |
|
return f"Error processing task {task['id']}: {e}" |
|
|
|
return "Successed! No pending tasks found." |
|
|
|
|
|
css = """ |
|
#col-container { |
|
margin: 0 auto; |
|
max-width: 800px; |
|
} |
|
""" |
|
|
|
with gr.Blocks(css=css) as app: |
|
with gr.Column(elem_id="col-container"): |
|
gr.Markdown(f"""# Tag The Image |
|
Get tag based on images using the Florence-2-base-PromptGen-v1.5 model. |
|
""") |
|
|
|
run_button = gr.Button("Submit", scale=0, elem_id="run-button") |
|
result = gr.Textbox(label="Generated Text", show_label=False) |
|
|
|
gr.on( |
|
triggers=[run_button.click], |
|
fn=infer, |
|
inputs=[], |
|
outputs=[result] |
|
) |
|
|
|
app.queue() |
|
app.launch(show_error=True) |
|
|
