File size: 2,398 Bytes
599abe6 |
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 |
from PIL import Image
import io
from pathlib import Path
import time
import traceback
from typing import Union, Optional
from .api_utils import upload_to_freeimage
def prepare_image(image: Union[str, bytes, Image.Image, None], status_tracker) -> Optional[str]:
"""
Prepare an image for use with LumaAI by resizing and uploading to CDN.
Args:
image: Input image (can be path, bytes, or PIL Image)
status_tracker: StatusTracker instance for progress updates
Returns:
Optional[str]: CDN URL of the prepared image, or None if no image provided
Raises:
Exception: If image preparation fails
"""
if image is None:
return None
try:
status_tracker.update_message("Preparing your image for its big moment...", 0.01)
# Convert to PIL Image if needed
if isinstance(image, str):
image = Image.open(image)
elif isinstance(image, bytes):
image = Image.open(io.BytesIO(image))
elif not isinstance(image, Image.Image):
raise Exception("That doesn't look like an image (unless I need glasses)")
# Resize image to 512x512
image = image.resize((512, 512), Image.Resampling.LANCZOS)
status_tracker.add_step("Image resized to 512x512", 0.02)
# Convert to RGB if necessary
if image.mode not in ('RGB', 'RGBA'):
image = image.convert('RGB')
# Save to temporary file
temp_dir = Path("temp")
temp_dir.mkdir(exist_ok=True)
temp_path = temp_dir / f"temp_image_{int(time.time())}.png"
image.save(str(temp_path), format='PNG', optimize=True)
# Upload to freeimage.host
try:
cdn_url = upload_to_freeimage(temp_path, status_tracker)
# Clean up temporary file
if temp_path.exists():
temp_path.unlink()
return cdn_url
except Exception as e:
# Clean up temporary file in case of error
if temp_path.exists():
temp_path.unlink()
raise Exception(f"Image upload failed: {str(e)}")
except Exception as e:
print(f"Error preparing image: {str(e)}")
print(traceback.format_exc())
raise Exception(f"Image preparation failed: {str(e)}")
|