Spaces:
Running
Running
| import requests | |
| import base64 | |
| import json | |
| 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() | |
| 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) | |
| # 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 and optionally downloads the result. | |
| 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, 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)) | |
| response.raise_for_status() | |
| api_response_data = response.json() | |
| # --- Optional Download Logic --- | |
| if download_path and isinstance(api_response_data, dict): | |
| image_url = api_response_data.get("imageUrl") | |
| if image_url: | |
| _download_image_from_url(image_url, download_path) | |
| else: | |
| print("Warning: 'imageUrl' not found in response, could not download image.") | |
| return api_response_data | |
| except requests.exceptions.RequestException as e: | |
| print(f"API request failed: {e}") | |
| return None |