297 lines
8.8 KiB
Python
297 lines
8.8 KiB
Python
import json
|
|
import folder_paths
|
|
import os
|
|
import requests
|
|
import base64
|
|
import traceback
|
|
import mimetypes
|
|
from PIL import Image, ImageSequence, ImageOps
|
|
import numpy as np
|
|
import imghdr
|
|
import mimetypes
|
|
import subprocess
|
|
import torch
|
|
import node_helpers
|
|
from matplotlib import font_manager
|
|
|
|
chunk_size = 5 * 1024 * 1024 # 1MB chunks
|
|
|
|
|
|
def get_system_font_files():
|
|
font_files = []
|
|
for font in font_manager.fontManager.ttflist:
|
|
font_file = os.path.basename(font.fname)
|
|
font_files.append(font_file)
|
|
return font_files
|
|
|
|
|
|
def is_image(file_path):
|
|
# Check if it's a common image type
|
|
if imghdr.what(file_path) is not None:
|
|
return True
|
|
# Additional check for SVG files
|
|
mime_type, _ = mimetypes.guess_type(file_path)
|
|
return mime_type is not None and mime_type.startswith("image")
|
|
|
|
|
|
def is_video(file_path):
|
|
video_extensions = [".mp4", ".avi", ".mov", ".mkv", ".flv", ".wmv"]
|
|
_, extension = os.path.splitext(file_path.lower())
|
|
mime_type, _ = mimetypes.guess_type(file_path)
|
|
return extension in video_extensions or (
|
|
mime_type is not None and mime_type.startswith("video")
|
|
)
|
|
|
|
|
|
def is_gif(file_path):
|
|
return imghdr.what(file_path) == "gif"
|
|
|
|
|
|
def get_video_duration(file_path):
|
|
if not is_video(file_path):
|
|
return None
|
|
|
|
try:
|
|
result = subprocess.run(
|
|
[
|
|
"ffprobe",
|
|
"-v",
|
|
"quiet",
|
|
"-print_format",
|
|
"json",
|
|
"-show_format",
|
|
"-show_streams",
|
|
file_path,
|
|
],
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
|
|
data = json.loads(result.stdout)
|
|
duration = float(data["format"]["duration"])
|
|
return duration
|
|
except (subprocess.SubprocessError, KeyError, json.JSONDecodeError):
|
|
return None
|
|
|
|
|
|
def images_file_to_tensor(image):
|
|
image_path = folder_paths.get_annotated_filepath(image)
|
|
img = node_helpers.pillow(Image.open, image_path)
|
|
return images_data_to_tensor(img)
|
|
|
|
|
|
def images_data_to_tensor(img):
|
|
|
|
output_images = []
|
|
output_masks = []
|
|
w, h = None, None
|
|
|
|
excluded_formats = ["MPO"]
|
|
|
|
for i in ImageSequence.Iterator(img):
|
|
i = node_helpers.pillow(ImageOps.exif_transpose, i)
|
|
|
|
if i.mode == "I":
|
|
i = i.point(lambda i: i * (1 / 255))
|
|
image = i.convert("RGB")
|
|
|
|
if len(output_images) == 0:
|
|
w = image.size[0]
|
|
h = image.size[1]
|
|
|
|
if image.size[0] != w or image.size[1] != h:
|
|
continue
|
|
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,]
|
|
if "A" in i.getbands():
|
|
mask = np.array(i.getchannel("A")).astype(np.float32) / 255.0
|
|
mask = 1.0 - torch.from_numpy(mask)
|
|
else:
|
|
mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
|
|
output_images.append(image)
|
|
output_masks.append(mask.unsqueeze(0))
|
|
|
|
if len(output_images) > 1 and img.format not in excluded_formats:
|
|
output_image = torch.cat(output_images, dim=0)
|
|
output_mask = torch.cat(output_masks, dim=0)
|
|
else:
|
|
output_image = output_images[0]
|
|
output_mask = output_masks[0]
|
|
|
|
return (output_image, output_mask)
|
|
|
|
|
|
def images_tensor_to_file(images, output_dir, compress_level, extension="png"):
|
|
filename_prefix = "socialman"
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = (
|
|
folder_paths.get_save_image_path(
|
|
filename_prefix, output_dir, images[0].shape[1], images[0].shape[0]
|
|
)
|
|
)
|
|
results = list()
|
|
for batch_number, image in enumerate(images):
|
|
i = 255.0 * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
|
|
filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
|
|
file = f"{filename_with_batch_num}_{counter:05}_.{extension}"
|
|
img.save(
|
|
os.path.join(full_output_folder, file),
|
|
compress_level=compress_level,
|
|
)
|
|
results.append(f"{folder_paths.get_output_directory()}/{file}")
|
|
counter += 1
|
|
return results
|
|
|
|
|
|
def get_file_base64(file_path):
|
|
with open(file_path, "rb") as file:
|
|
content = file.read()
|
|
file_content_base64 = base64.b64encode(content).decode("utf-8")
|
|
return file_content_base64
|
|
|
|
|
|
def upload_file_to_signed_s3(file_path, presigned_url):
|
|
# Check if file exists
|
|
if not os.path.isfile(file_path):
|
|
raise FileNotFoundError(f"File not found: {file_path}")
|
|
|
|
content_type, _ = mimetypes.guess_type(file_path)
|
|
if content_type is None:
|
|
content_type = "application/octet-stream"
|
|
# Get file size
|
|
file_size = os.path.getsize(file_path)
|
|
|
|
# Open file in binary mode
|
|
with open(file_path, "rb") as file:
|
|
# Use requests to PUT the file to the pre-signed URL
|
|
response = requests.put(
|
|
presigned_url,
|
|
data=file,
|
|
headers={"Content-Length": str(file_size), "Content-Type": content_type},
|
|
)
|
|
|
|
# Check if the upload was successful
|
|
if response.status_code == 200:
|
|
print(f"File {file_path} uploaded successfully.")
|
|
else:
|
|
print(f"Failed to upload file. Status code: {response.status_code}")
|
|
print(f"Response: {response.text}")
|
|
|
|
|
|
def upload_file(file_path, api_base_url, auth_token):
|
|
try:
|
|
# Initiate upload
|
|
headers = {"Authorization": auth_token}
|
|
total_chunks = calculate_total_chunks(file_path)
|
|
|
|
print(f"Initiating upload for file: {file_path}")
|
|
init_response = requests.post(
|
|
f"{api_base_url}/initiate-upload",
|
|
headers=headers,
|
|
json={
|
|
"fileName": os.path.basename(file_path),
|
|
"totalChunks": total_chunks,
|
|
},
|
|
)
|
|
init_response.raise_for_status()
|
|
print(f"Initiation response: {init_response.text}")
|
|
|
|
upload_id = init_response.json()["uploadId"]
|
|
print(f"Upload ID: {upload_id}")
|
|
|
|
# Read file in chunks and upload
|
|
|
|
chunk_number = 1
|
|
|
|
with open(file_path, "rb") as f:
|
|
while chunk_number <= total_chunks:
|
|
chunk = f.read(chunk_size)
|
|
if not chunk:
|
|
break
|
|
|
|
upload_url = f"{api_base_url}/upload-chunk/{upload_id}/{chunk_number}"
|
|
print(f"upload_url: {upload_url}")
|
|
print(f"Upload ID: {upload_id}")
|
|
print(f"Uploading chunk {chunk_number}/{total_chunks}")
|
|
response = requests.put(
|
|
upload_url, headers=headers, files={"file": chunk}
|
|
)
|
|
# response.raise_for_status()
|
|
print(f"Chunk {chunk_number} upload response: {response.text}")
|
|
|
|
chunk_number += 1
|
|
|
|
# Complete upload
|
|
complete_url = f"{api_base_url}/complete-upload/{upload_id}"
|
|
print("Completing upload")
|
|
complete_response = requests.post(complete_url, headers=headers)
|
|
complete_response.raise_for_status()
|
|
print(f"Complete upload response: {complete_response.text}")
|
|
|
|
return upload_id
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
print(f"Request failed: {e}")
|
|
print(
|
|
f"Response content: {e.response.content if e.response else 'No response'}"
|
|
)
|
|
print(f"Traceback: {traceback.format_exc()}")
|
|
raise
|
|
|
|
except Exception as e:
|
|
print(f"Unexpected error: {e}")
|
|
print(f"Traceback: {traceback.format_exc()}")
|
|
raise
|
|
|
|
|
|
def calculate_total_chunks(file_path):
|
|
file_size = os.path.getsize(file_path)
|
|
return -(-file_size // chunk_size) # Ceiling division
|
|
|
|
|
|
def image_files_only():
|
|
image_extensions = (".jpg", ".jpeg", ".png", ".gif", ".bmp", ".tiff", ".webp")
|
|
input_dir = folder_paths.get_input_directory()
|
|
return [
|
|
f
|
|
for f in os.listdir(input_dir)
|
|
if os.path.isfile(os.path.join(input_dir, f))
|
|
and f.lower().endswith(image_extensions)
|
|
]
|
|
|
|
|
|
def mask_string(input_string):
|
|
# Ensure the input is a string
|
|
input_string = str(input_string)
|
|
|
|
# If the string is 5 characters or longer
|
|
if len(input_string) >= 5:
|
|
return "***" + input_string[-5:]
|
|
# If the string is shorter than 5 characters
|
|
else:
|
|
return "***" + input_string
|
|
|
|
|
|
def write_json_to_file(filename, data):
|
|
with open(filename, "w") as file:
|
|
json.dump(data, file, indent=4)
|
|
|
|
|
|
def read_json_from_file(filename):
|
|
try:
|
|
with open(filename, "r") as file:
|
|
return json.load(file)
|
|
except FileNotFoundError:
|
|
return {"error": "File not found."}
|
|
except json.JSONDecodeError:
|
|
return {"error": "Invalid JSON in file."}
|
|
|
|
|
|
def update_json_file(filename, new_data):
|
|
old_data = read_json_from_file(filename)
|
|
old_data.update(new_data)
|
|
with open(filename, "w") as file:
|
|
json.dump(old_data, file, indent=4)
|