from aiohttp import web import os import json import folder_paths from PIL import Image import server # Import server for node_info import uuid # For generating unique filenames async def get_hello(request: web.Request) -> web.Response: return web.json_response({"status": "success", "message": "Hello from the CozyGen API!"}) async def get_gallery_files(request: web.Request) -> web.Response: subfolder = request.rel_url.query.get('subfolder', '') try: page = int(request.rel_url.query.get('page', '1')) per_page = int(request.rel_url.query.get('per_page', '20')) except ValueError: return web.json_response({"error": "Invalid page or per_page parameter"}, status=400) output_directory = folder_paths.get_output_directory() # Security: Prevent directory traversal gallery_path = os.path.normpath(os.path.join(output_directory, subfolder)) if not gallery_path.startswith(output_directory): return web.json_response({"error": "Unauthorized path"}, status=403) if not os.path.exists(gallery_path) or not os.path.isdir(gallery_path): return web.json_response({"error": "Gallery directory not found"}, status=404) items = os.listdir(gallery_path) gallery_items = [] for item_name in items: item_path = os.path.join(gallery_path, item_name) if os.path.isdir(item_path): mod_time = os.path.getmtime(item_path) gallery_items.append({ "filename": item_name, "type": "directory", "subfolder": os.path.join(subfolder, item_name), "mod_time": mod_time }) elif item_name.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.webp', '.mp4', '.webm', '.mp3', '.wav', '.flac')): mod_time = os.path.getmtime(item_path) gallery_items.append({ "filename": item_name, "type": "output", "subfolder": subfolder, "mod_time": mod_time }) # Sort items: directories first, then by modification time gallery_items.sort(key=lambda x: (x['type'] == 'directory', x.get('mod_time', 0)), reverse=True) # Pagination total_items = len(gallery_items) total_pages = (total_items + per_page - 1) // per_page start_index = (page - 1) * per_page end_index = start_index + per_page paginated_items = gallery_items[start_index:end_index] # Remove mod_time before sending for item in paginated_items: if 'mod_time' in item: del item['mod_time'] return web.json_response({ "items": paginated_items, "page": page, "per_page": per_page, "total_pages": total_pages, "total_items": total_items }) async def upload_image(request: web.Request) -> web.Response: reader = await request.multipart() field = await reader.next() if field.name != 'image': return web.json_response({"error": "Expected field 'image'"}, status=400) filename = field.filename if not filename: return web.json_response({"error": "No filename provided"}, status=400) # Generate a unique filename to prevent collisions unique_filename = f"{uuid.uuid4()}_{filename}" # Save to the input directory of ComfyUI input_dir = folder_paths.get_input_directory() save_path = os.path.join(input_dir, unique_filename) size = 0 with open(save_path, 'wb') as f: while True: chunk = await field.read_chunk() if not chunk: break f.write(chunk) size += len(chunk) return web.json_response({"filename": unique_filename, "size": size}) async def get_workflow_list(request: web.Request) -> web.Response: workflows_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "workflows") if not os.path.exists(workflows_dir): return web.json_response({"error": "Workflows directory not found"}, status=404) workflow_files = [f for f in os.listdir(workflows_dir) if f.endswith('.json')] return web.json_response({"workflows": workflow_files}) async def get_workflow_file(request: web.Request) -> web.Response: filename = request.match_info.get('filename', '') workflows_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "workflows") workflow_path = os.path.join(workflows_dir, filename) if not os.path.exists(workflow_path): return web.json_response({"error": f"Workflow file '{filename}' not found"}, status=404) try: with open(workflow_path, 'r', encoding='utf-8') as f: workflow_content = json.load(f) return web.json_response(workflow_content) except json.JSONDecodeError: return web.json_response({"error": f"Invalid JSON in workflow file '{filename}'"}, status=400) except Exception as e: return web.json_response({"error": f"Error reading workflow file: {e}"}, status=500) import comfy.samplers # Get all valid model folder types from ComfyUI itself valid_model_types = folder_paths.folder_names_and_paths.keys() # A map for aliases to official folder_paths names alias_map = { "unet": "unet_gguf" } async def get_choices(request: web.Request) -> web.Response: choice_type = request.rel_url.query.get('type', '') if not choice_type: return web.json_response({"error": "Missing 'type' query parameter"}, status=400) # Alias map for backward compatibility alias_map = { "samplers_list": "sampler", "schedulers_list": "scheduler", "unet": "unet_gguf" # Example of another potential alias } resolved_choice_type = alias_map.get(choice_type, choice_type) choices = [] if resolved_choice_type == "scheduler": choices = comfy.samplers.KSampler.SCHEDULERS elif resolved_choice_type == "sampler": choices = comfy.samplers.KSampler.SAMPLERS elif resolved_choice_type in valid_model_types: choices = folder_paths.get_filename_list(resolved_choice_type) else: return web.json_response({"error": f"Invalid choice type: {choice_type}"}, status=400) return web.json_response({"choices": choices}) routes = [ web.get('/cozygen/hello', get_hello), web.get('/cozygen/gallery', get_gallery_files), web.post('/cozygen/upload_image', upload_image), web.get('/cozygen/workflows', get_workflow_list), web.get('/cozygen/workflows/{filename}', get_workflow_file), web.get('/cozygen/get_choices', get_choices), ]