Spaces:
Sleeping
Sleeping
File size: 4,707 Bytes
4327cda |
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 |
import os
from typing import Optional
from PIL import Image
import io
import base64
import requests
from dotenv import load_dotenv
load_dotenv()
class ImageDryer:
def __init__(self):
"""Initialize the ImageDryer with Stability AI API."""
self.api_key = os.getenv("STABILITY_API_KEY")
self.api_host = "https://api.stability.ai"
self.engine_id = "stable-diffusion-xl-1024-v1-0"
def preprocess_image(self, image: Image.Image) -> Image.Image:
"""Preprocess the image to meet API requirements."""
try:
# Convert to RGB if needed
if image.mode != "RGB":
image = image.convert("RGB")
# Get current dimensions
width, height = image.size
# Calculate aspect ratio
aspect_ratio = width / height
# Define supported dimensions
supported_dimensions = [
(1024, 1024), # 1:1
(1152, 896), # 1.29:1
(1216, 832), # 1.46:1
(1344, 768), # 1.75:1
(1536, 640), # 2.4:1
(640, 1536), # 1:2.4
(768, 1344), # 1:1.75
(832, 1216), # 1:1.46
(896, 1152) # 1:1.29
]
# Find the closest aspect ratio
target_dims = min(supported_dimensions,
key=lambda dims: abs((dims[0]/dims[1]) - aspect_ratio))
# Print debug info
print(f"Original dimensions: {width}x{height}, aspect ratio: {aspect_ratio:.2f}")
print(f"Selected target dimensions: {target_dims[0]}x{target_dims[1]}")
# Resize image to target dimensions using high-quality resampling
image = image.resize(target_dims, Image.Resampling.LANCZOS)
# Verify final dimensions
final_width, final_height = image.size
if (final_width, final_height) not in supported_dimensions:
raise ValueError(f"Failed to resize to supported dimensions. Got {final_width}x{final_height}")
return image
except Exception as e:
print(f"Error in preprocess_image: {str(e)}")
raise
def process_image(self, image: Image.Image) -> Optional[Image.Image]:
"""Process an image to make it appear dry using Stability AI API."""
try:
# Preprocess the image
processed_image = self.preprocess_image(image)
# Convert image to bytes
buffered = io.BytesIO()
processed_image.save(buffered, format="PNG")
img_bytes = buffered.getvalue()
# Prepare the API request
url = f"{self.api_host}/v1/generation/{self.engine_id}/image-to-image"
headers = {
"Authorization": f"Bearer {self.api_key}"
}
files = {
"init_image": ("image.png", img_bytes, "image/png"),
}
data = {
"image_strength": 0.35,
"text_prompts[0][text]": "A completely dry version of this item, photorealistic, detailed texture, no water or moisture",
"text_prompts[0][weight]": 1,
"text_prompts[1][text]": "wet, moist, damp, water droplets, puddles, stains",
"text_prompts[1][weight]": -1,
"cfg_scale": 7,
"samples": 1,
"steps": 30
}
# Make the API request
response = requests.post(url, headers=headers, files=files, data=data)
if response.status_code != 200:
raise Exception(f"API request failed: {response.text}")
# Process the response
data = response.json()
image_data = base64.b64decode(data["artifacts"][0]["base64"])
# Convert to PIL Image
result = Image.open(io.BytesIO(image_data))
return result
except Exception as e:
print(f"Error processing image: {str(e)}")
return None
def save_image(self, image: Image.Image, filename: str) -> None:
"""Save an image to a file."""
image.save(filename)
def image_to_bytes(self, image: Image.Image) -> bytes:
"""Convert PIL Image to bytes."""
img_byte_arr = io.BytesIO()
image.save(img_byte_arr, format='PNG')
return img_byte_arr.getvalue() |