Spaces:
Paused
Paused
File size: 5,806 Bytes
fceeb2f |
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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
import base64
import io
import json
import os
import time
from typing import Any, Dict, Optional
from PIL import Image
import requests
def _image_to_base64(image: Image.Image) -> str:
buffer = io.BytesIO()
image_format = (image.format or "PNG").upper()
if image_format not in {"PNG", "JPEG", "JPG"}:
image_format = "PNG"
image.save(buffer, format=image_format)
return base64.b64encode(buffer.getvalue()).decode("utf-8")
def _extract_status(payload: Dict[str, Any]) -> Optional[str]:
status_info = payload.get("status") or payload.get("state")
if isinstance(status_info, dict):
state = status_info.get("state") or status_info.get("status")
if isinstance(state, str):
return state.lower()
elif isinstance(status_info, str):
return status_info.lower()
return None
def _poll_bria_status(
status_url: str,
headers: Dict[str, str],
timeout_seconds: int = 120,
poll_interval: float = 1.5,
) -> Dict[str, Any]:
deadline = time.time() + timeout_seconds
while True:
response = requests.get(status_url, headers=headers, timeout=30)
response.raise_for_status()
payload: Dict[str, Any] = response.json()
state = _extract_status(payload)
if state in {"succeeded", "success", "completed", "done"}:
if isinstance(payload.get("result"), dict):
return payload["result"]
if payload.get("results") is not None:
return payload["results"]
return payload
if state in {"failed", "error", "cancelled", "canceled"}:
raise RuntimeError(
f"Bria VLM API request failed: {json.dumps(payload, indent=2)}"
)
if time.time() > deadline:
raise TimeoutError(
f"Bria VLM API request timed out while polling {status_url}"
)
time.sleep(poll_interval)
def _submit_bria_request(
url: str, payload: Dict[str, Any], api_token: str
) -> Dict[str, Any]:
headers = {
"Content-Type": "application/json",
"api_token": api_token,
}
response = requests.post(url, json=payload, headers=headers, timeout=30)
response.raise_for_status()
initial_payload: Dict[str, Any] = response.json()
status_url = (
initial_payload.get("status_url")
or initial_payload.get("statusUrl")
or (initial_payload.get("status") or {}).get("status_url")
)
if status_url:
return _poll_bria_status(status_url, headers)
if isinstance(initial_payload.get("result"), dict):
return initial_payload["result"]
if initial_payload.get("results") is not None:
return initial_payload["results"]
return initial_payload
def _parse_vlm_response(data: Any, prompt_role: str) -> str:
if isinstance(data, dict):
direct_match = data.get(prompt_role)
if isinstance(direct_match, str):
return direct_match
for key in ("prompt", "structured_prompt", "structuredPrompt", "text"):
if key in data:
value = data[key]
if isinstance(value, str):
return value
if isinstance(value, dict):
nested = value.get(prompt_role)
if isinstance(nested, str):
return nested
for key in ("result", "results"):
if key in data:
nested_result = _parse_vlm_response(data[key], prompt_role)
if nested_result:
return nested_result
if isinstance(data, list):
for item in data:
nested_result = _parse_vlm_response(item, prompt_role)
if nested_result:
return nested_result
return json.dumps(data)
def get_prompt_api(image_path: str, prompt_role: str) -> str:
"""Send an image to the Bria VLM API and return the extracted prompt text.
The payload keys are aligned with the current public docs but may require
adjustment if your Bria workspace is configured differently. Override the
default endpoint via the ``BRIA_API_VLM_ENDPOINT`` environment variable if
you are using a custom workflow.
"""
api_token = os.environ.get("BRIA_API_KEY")
if not api_token:
raise EnvironmentError(
"BRIA_API_KEY environment variable is required to use the Bria VLM API."
)
base_url = os.environ.get("BRIA_API_BASE_URL", "https://engine.prod.bria-api.com")
endpoint = os.environ.get("BRIA_API_VLM_ENDPOINT", "/v2/structured_prompt/generate")
url = f"{base_url.rstrip('/')}{endpoint}"
# convert image to base64
with Image.open(image_path) as image:
image_b64 = _image_to_base64(image)
payload = {"images": [image_b64]}
response = _submit_bria_request(url, payload, api_token)
return response["structured_prompt"]
def get_image_from_url(image_url: str) -> Image.Image:
"""Get an image from a URL."""
response = requests.get(image_url)
return Image.open(io.BytesIO(response.content))
def generate_image(prompt: str) -> Image.Image:
"""Generate an image from a prompt using the Bria VLM API."""
api_token = os.environ.get("BRIA_API_KEY")
if not api_token:
raise EnvironmentError(
"BRIA_API_KEY environment variable is required to use the Bria VLM API."
)
base_url = os.environ.get("BRIA_API_BASE_URL", "https://engine.prod.bria-api.com")
endpoint = os.environ.get("BRIA_API_GENERATE_ENDPOINT", "/v2/image/generate")
url = f"{base_url.rstrip('/')}{endpoint}"
payload = {"structured_prompt": prompt}
response = _submit_bria_request(url, payload, api_token)
return get_image_from_url(response["image_url"])
|