Spaces:
Running
Running
File size: 6,013 Bytes
c331dad 40babb7 c331dad 40babb7 c331dad 40babb7 c331dad 40babb7 c331dad 40babb7 c331dad 40babb7 c331dad 40babb7 c331dad 40babb7 c331dad 40babb7 c331dad 40babb7 |
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 |
import requests
import base64
import json
import os
from io import BytesIO
from typing import Optional, Dict
# Pillow is required for image format conversion and normalization.
# Please install it using: pip install Pillow
try:
from PIL import Image
except ImportError:
print("Pillow library not found. Please install it using: pip install Pillow")
exit()
# --- Configuration ---
IMGBB_API_KEY = os.getenv("IMGBB_API_KEY")
def upload_to_imgbb(image_path: str, file_name: str) -> Optional[str]:
"""
Uploads the image at image_path to ImgBB.
Returns the public URL or None on failure.
"""
if not IMGBB_API_KEY:
print("Warning: IMGBB_API_KEY not set, skipping upload")
return None
print(f"Uploading {file_name} to ImgBB...")
try:
with open(image_path, 'rb') as f:
files = {"image": (file_name, f.read())}
resp = requests.post(
"https://api.imgbb.com/1/upload",
params={"key": IMGBB_API_KEY},
files=files,
timeout=20
)
resp.raise_for_status()
data = resp.json().get("data", {})
url = data.get("url")
if url:
print(f"Successfully uploaded to ImgBB: {url}")
return url
else:
print(f"Error: ImgBB API response missing 'url'. Response: {resp.json()}")
return None
except requests.exceptions.RequestException as e:
print(f"Error: ImgBB upload failed: {e}")
return None
def _download_image_from_url(image_url: str, save_path: str) -> bool:
"""
Downloads an image from a URL and saves it to a local path.
Args:
image_url: The URL of the image to download.
save_path: The local path to save the downloaded image.
Returns:
True if the download was successful, False otherwise.
"""
print(f"Downloading generated image from: {image_url}")
try:
image_response = requests.get(image_url, stream=True, timeout=30)
# Check if the download request was successful.
if image_response.status_code == 200:
content_type = image_response.headers.get('Content-Type', '')
if 'image' in content_type:
with open(save_path, 'wb') as f:
for chunk in image_response.iter_content(1024):
f.write(chunk)
print(f"Image successfully saved to {save_path}")
return True
else:
print(f"Error: Content at URL is not an image. Content-Type: {content_type}")
return False
else:
print(f"Error: Failed to download image. Status code: {image_response.status_code}")
return False
except requests.exceptions.RequestException as e:
print(f"An error occurred during image download: {e}")
return False
def generate_image(
prompt_text: str,
image_path: str,
download_path: Optional[str] = None
) -> Optional[Dict]:
"""
Sends a request to the image generation API, optionally downloads the result,
and uploads it to ImgBB.
Args:
prompt_text: The instructional text for image modification.
image_path: The file path to the input image (any common format).
download_path: If provided, the path to save the generated image.
Returns:
A dictionary of the JSON response from the API, including the ImgBB URL,
or None on error.
"""
url = "https://kontext-chat.replicate.dev/generate-image"
try:
# --- Image Normalization Step ---
with Image.open(image_path) as img:
if img.mode != 'RGB':
img = img.convert('RGB')
with BytesIO() as output_buffer:
img.save(
output_buffer,
format="JPEG",
quality=95,
subsampling=0,
progressive=False
)
image_bytes = output_buffer.getvalue()
except FileNotFoundError:
print(f"Error: Image file not found at {image_path}")
return None
except Exception as e:
print(f"Error processing image file. Ensure it's a valid image. Details: {e}")
return None
encoded_string = base64.b64encode(image_bytes).decode('utf-8')
input_image_data_uri = f"data:image/jpeg;base64,{encoded_string}"
payload = {
"prompt": prompt_text,
"input_image": input_image_data_uri
}
headers = {
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:139.0) Gecko/20100101 Firefox/139.0",
"Accept": "*/*",
"Content-Type": "application/json",
"Referer": "https://kontext-chat.replicate.dev/",
"Origin": "https://kontext-chat.replicate.dev",
}
try:
response = requests.post(url, headers=headers, data=json.dumps(payload), timeout=45)
response.raise_for_status()
api_response_data = response.json()
imgbb_url = None # Initialize imgbb_url
# --- Download & Upload Logic ---
if download_path and isinstance(api_response_data, dict):
image_url = api_response_data.get("imageUrl")
if image_url:
if _download_image_from_url(image_url, download_path):
# If download is successful, upload the saved file to ImgBB
file_name = os.path.basename(download_path)
imgbb_url = upload_to_imgbb(download_path, file_name)
else:
print("Warning: 'imageUrl' not found in API response, cannot download or upload image.")
# Add the ImgBB URL to the response dictionary if it was generated
if isinstance(api_response_data, dict):
api_response_data['imgbb_url'] = imgbb_url
return api_response_data
except requests.exceptions.RequestException as e:
print(f"API request failed: {e}")
return None
|