Spaces:
Sleeping
Sleeping
File size: 4,001 Bytes
7ff46e3 99b7a8d 7ff46e3 99b7a8d | 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 | 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
|