Spaces:
Sleeping
Sleeping
| import random | |
| import tempfile | |
| import zipfile | |
| import os | |
| import uuid | |
| import PIL.Image | |
| import PIL.ImageOps | |
| import gradio as gr | |
| from settings import MAX_SEED | |
| def randomize_seed_fn(seed: int, randomize_seed: bool) -> int: | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) # noqa: S311 | |
| return seed | |
| def wrap_process_with_zip(process_fn, has_color_correction=True): | |
| def wrapped(*args): | |
| try: | |
| if has_color_correction: | |
| # The last argument of args is the color_correction value | |
| color_corr_val = args[-1] | |
| # The remaining arguments are passed to the original process_fn | |
| actual_args = args[:-1] | |
| else: | |
| color_corr_val = None | |
| actual_args = args | |
| res_images = process_fn(*actual_args) | |
| except Exception as e: | |
| print(f"Error in process: {e}") | |
| raise e | |
| if not res_images: | |
| return [], gr.update(visible=False) | |
| processed_res = [] | |
| for img in res_images: | |
| if color_corr_val == "white bg, black lines": | |
| # Invert the image (black bg -> white bg, white lines -> black lines) | |
| if isinstance(img, PIL.Image.Image): | |
| img_rgb = img.convert("RGB") | |
| inverted_img = PIL.ImageOps.invert(img_rgb) | |
| processed_res.append(inverted_img) | |
| elif isinstance(img, str) and os.path.exists(img): | |
| try: | |
| img_pil = PIL.Image.open(img).convert("RGB") | |
| inverted_img = PIL.ImageOps.invert(img_pil) | |
| processed_res.append(inverted_img) | |
| except Exception: | |
| processed_res.append(img) | |
| else: | |
| # Try to convert numpy array or other format | |
| try: | |
| img_pil = PIL.Image.fromarray(img).convert("RGB") | |
| inverted_img = PIL.ImageOps.invert(img_pil) | |
| processed_res.append(inverted_img) | |
| except Exception: | |
| processed_res.append(img) | |
| else: | |
| processed_res.append(img) | |
| res_images = processed_res | |
| # Now create the ZIP file from res_images | |
| temp_dir = tempfile.gettempdir() | |
| zip_filename = os.path.join(temp_dir, f"controlnet_results_{uuid.uuid4().hex[:8]}.zip") | |
| try: | |
| with zipfile.ZipFile(zip_filename, 'w') as zipf: | |
| for idx, img in enumerate(res_images): | |
| if isinstance(img, PIL.Image.Image): | |
| temp_img_path = os.path.join(temp_dir, f"temp_{uuid.uuid4().hex[:8]}.png") | |
| img.save(temp_img_path, "PNG") | |
| zipf.write(temp_img_path, f"result_{idx}.png") | |
| try: | |
| os.remove(temp_img_path) | |
| except Exception: | |
| pass | |
| elif isinstance(img, str) and os.path.exists(img): | |
| zipf.write(img, os.path.basename(img)) | |
| else: | |
| try: | |
| temp_img_path = os.path.join(temp_dir, f"temp_{uuid.uuid4().hex[:8]}.png") | |
| PIL.Image.fromarray(img).save(temp_img_path, "PNG") | |
| zipf.write(temp_img_path, f"result_{idx}.png") | |
| try: | |
| os.remove(temp_img_path) | |
| except Exception: | |
| pass | |
| except Exception: | |
| pass | |
| return res_images, gr.update(value=zip_filename, visible=True) | |
| except Exception as zip_err: | |
| print(f"Error creating zip file: {zip_err}") | |
| return res_images, gr.update(visible=False) | |
| return wrapped | |