Spaces:
Sleeping
Sleeping
| import os | |
| import gradio as gr | |
| import requests | |
| import base64 | |
| from PIL import Image | |
| from io import BytesIO | |
| import tempfile | |
| import zipfile | |
| API_KEY = os.getenv("SEGMIND_API_KEY") | |
| API_URL = "https://api.segmind.com/v1/nomos-upscaler" | |
| def image_file_to_base64(image_path): | |
| with open(image_path, 'rb') as f: | |
| image_data = f.read() | |
| return base64.b64encode(image_data).decode('utf-8') | |
| def image_url_to_base64(image_url): | |
| response = requests.get(image_url) | |
| image_data = response.content | |
| return base64.b64encode(image_data).decode('utf-8') | |
| def upscale_single_image(base64_image): | |
| data = { | |
| "image": base64_image, | |
| "image_format": "png", | |
| "image_quality": 95, | |
| "base64": False | |
| } | |
| headers = {'x-api-key': API_KEY} | |
| response = requests.post(API_URL, json=data, headers=headers) | |
| if response.status_code == 200: | |
| return Image.open(BytesIO(response.content)) | |
| else: | |
| raise Exception(f"API Error: {response.status_code} - {response.text}") | |
| def upscale_images(files, urls, state): | |
| results = [] | |
| filepaths = [] | |
| temp_dir = tempfile.mkdtemp() | |
| # Process uploaded files | |
| if files: | |
| for file_path in files: | |
| try: | |
| base64_img = image_file_to_base64(file_path) | |
| upscaled_img = upscale_single_image(base64_img) | |
| out_path = os.path.join(temp_dir, f"upscaled_{os.path.basename(file_path)}.png") | |
| upscaled_img.save(out_path, format="PNG") | |
| results.append(upscaled_img) | |
| filepaths.append(out_path) | |
| except Exception as e: | |
| results.append(f"Error: {str(e)}") | |
| # Process URLs | |
| if urls: | |
| url_list = [u.strip() for u in urls.splitlines() if u.strip()] | |
| for idx, url in enumerate(url_list): | |
| try: | |
| base64_img = image_url_to_base64(url) | |
| upscaled_img = upscale_single_image(base64_img) | |
| out_path = os.path.join(temp_dir, f"upscaled_url_{idx+1}.png") | |
| upscaled_img.save(out_path, format="PNG") | |
| results.append(upscaled_img) | |
| filepaths.append(out_path) | |
| except Exception as e: | |
| results.append(f"Error: {str(e)}") | |
| # Update state with new filepaths | |
| state = filepaths | |
| return results, state | |
| def zip_upscaled_images(state): | |
| if not state: | |
| return None | |
| temp_zip = tempfile.NamedTemporaryFile(delete=False, suffix=".zip") | |
| with zipfile.ZipFile(temp_zip.name, 'w') as zipf: | |
| for fpath in state: | |
| zipf.write(fpath, arcname=os.path.basename(fpath)) | |
| return temp_zip.name | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Bulk Image Upscaler\nUpload multiple images or paste image URLs (one per line). Download all upscaled images as a zip.") | |
| state = gr.State([]) | |
| with gr.Row(): | |
| file_input = gr.File(label="Upload Images", file_count="multiple", type="filepath") | |
| url_input = gr.Textbox(label="Image URLs (one per line)", lines=5) | |
| output_gallery = gr.Gallery(label="Upscaled Images", show_label=True) | |
| run_btn = gr.Button("Upscale Images") | |
| download_zip_btn = gr.Button("Download All as ZIP") | |
| zip_file_output = gr.File(label="Download ZIP") | |
| run_btn.click( | |
| upscale_images, | |
| inputs=[file_input, url_input, state], | |
| outputs=[output_gallery, state] | |
| ) | |
| download_zip_btn.click( | |
| zip_upscaled_images, | |
| inputs=[state], | |
| outputs=zip_file_output | |
| ) | |
| if __name__ == "__main__": | |
| # Make sure to set the environment variable SEGMIND_API_KEY before running | |
| # For example, in Linux/Mac: export SEGMIND_API_KEY=your_api_key | |
| demo.launch() |