Spaces:
Runtime error
Runtime error
import gc | |
import os | |
import random | |
import numpy as np | |
import json | |
import uuid | |
from PIL import Image, PngImagePlugin | |
from datetime import datetime | |
from dataclasses import dataclass | |
from typing import Callable, Dict, Optional, Tuple | |
MAX_SEED = np.iinfo(np.int32).max | |
class StyleConfig: | |
prompt: str | |
negative_prompt: str | |
def randomize_seed_fn(seed: int, randomize_seed: bool) -> int: | |
if randomize_seed: | |
seed = random.randint(0, MAX_SEED) | |
return seed | |
def parse_aspect_ratio(aspect_ratio: str) -> Optional[Tuple[int, int]]: | |
if aspect_ratio == "Custom": | |
return None | |
width, height = aspect_ratio.split(" x ") | |
return int(width), int(height) | |
def aspect_ratio_handler( | |
aspect_ratio: str, custom_width: int, custom_height: int | |
) -> Tuple[int, int]: | |
if aspect_ratio == "Custom": | |
return custom_width, custom_height | |
else: | |
width, height = parse_aspect_ratio(aspect_ratio) | |
return width, height | |
def free_memory() -> None: | |
torch.cuda.empty_cache() | |
gc.collect() | |
def preprocess_prompt( | |
style_dict, | |
style_name: str, | |
positive: str, | |
negative: str = "", | |
add_style: bool = True, | |
) -> Tuple[str, str]: | |
p, n = style_dict.get(style_name, style_dict["(None)"]) | |
if add_style and positive.strip(): | |
formatted_positive = p.format(prompt=positive) | |
else: | |
formatted_positive = positive | |
combined_negative = n | |
if negative.strip(): | |
if combined_negative: | |
combined_negative += ", " + negative | |
else: | |
combined_negative = negative | |
return formatted_positive, combined_negative | |
def get_random_line_from_file(file_path: str) -> str: | |
with open(file_path, "r") as file: | |
lines = file.readlines() | |
if not lines: | |
return "" | |
return random.choice(lines).strip() | |
def preprocess_image_dimensions(width, height): | |
if width % 8 != 0: | |
width = width - (width % 8) | |
if height % 8 != 0: | |
height = height - (height % 8) | |
return width, height | |
from PIL import Image | |
import requests | |
from io import BytesIO | |
def save_image_replicate(image_url, metadata, output_dir, is_colab): | |
if colab: | |
current_time = datetime.now().strftime("%Y%m%d_%H%M%S") | |
filename = f"image_{current_time}.png" | |
else: | |
filename = str(uuid.uuid4()) + ".png" | |
os.makedirs(output_dir, exist_ok=True) | |
response = requests.get(image_url) | |
image = Image.open(BytesIO(response.content)) | |
filepath = os.path.join(output_dir, filename) | |
metadata_str = json.dumps(metadata) | |
info = PngImagePlugin.PngInfo() | |
info.add_text("metadata", metadata_str) | |
image.save(filepath, "PNG", pnginfo=info) | |
return filepath | |
def save_image(image, metadata, output_dir, is_colab): | |
if is_colab: | |
current_time = datetime.now().strftime("%Y%m%d_%H%M%S") | |
filename = f"image_{current_time}.png" | |
else: | |
filename = str(uuid.uuid4()) + ".png" | |
os.makedirs(output_dir, exist_ok=True) | |
filepath = os.path.join(output_dir, filename) | |
metadata_str = json.dumps(metadata) | |
info = PngImagePlugin.PngInfo() | |
info.add_text("metadata", metadata_str) | |
image.save(filepath, "PNG", pnginfo=info) | |
return filepath | |