import os import sys import json import time import requests import gradio as gr import numpy as np from PIL import Image import io import base64 import spaces # ComfyUI API endpoint COMFY_API = "http://127.0.0.1:8188/api" WORKFLOW_PATH = "/app/workflows/Workflow_12_11.json" # Load the workflow template try: with open(WORKFLOW_PATH, "r") as f: workflow_template = json.load(f) print(f"Loaded workflow template from {WORKFLOW_PATH}") except Exception as e: print(f"Error loading workflow template: {str(e)}") workflow_template = {} def queue_prompt(prompt): """Send a prompt to ComfyUI for processing""" p = {"prompt": prompt} try: response = requests.post(f"{COMFY_API}/prompt", json=p) return response.json() except Exception as e: print(f"Error queuing prompt: {str(e)}") return {"error": str(e)} def get_image(filename, subfolder, folder_type): """Get an image from ComfyUI's output folder""" try: response = requests.get(f"{COMFY_API}/view?filename={filename}&subfolder={subfolder}&type={folder_type}") return Image.open(io.BytesIO(response.content)) except Exception as e: print(f"Error getting image {filename}: {str(e)}") return None def upload_image(image, filename): """Upload an image to ComfyUI's input folder""" try: if isinstance(image, str): # Base64 string image_data = base64.b64decode(image.split(",")[1]) files = {"image": (filename, image_data)} else: # PIL Image or numpy array if isinstance(image, np.ndarray): image = Image.fromarray(image) img_byte_arr = io.BytesIO() image.save(img_byte_arr, format='PNG') img_byte_arr.seek(0) files = {"image": (filename, img_byte_arr.getvalue())} response = requests.post(f"{COMFY_API}/upload/image", files=files) return response.json() except Exception as e: print(f"Error uploading image: {str(e)}") return {"error": str(e)} def check_progress(prompt_id): """Check the progress of a ComfyUI prompt""" try: response = requests.get(f"{COMFY_API}/history/{prompt_id}") return response.json() except Exception as e: print(f"Error checking progress: {str(e)}") return {"error": str(e)} @spaces.GPU def generate_avatar(player_image, pose_image, shirt_image, player_name, team_name, age_gender_bg, shirt_color, style): """Generate a football player avatar using ComfyUI workflow""" # Upload images to ComfyUI player_upload = upload_image(player_image, "player.png") pose_upload = upload_image(pose_image, "pose.png") shirt_upload = upload_image(shirt_image, "shirt.png") if "error" in player_upload or "error" in pose_upload or "error" in shirt_upload: return None, f"Error uploading images: {player_upload.get('error', '')} {pose_upload.get('error', '')} {shirt_upload.get('error', '')}" # Create a copy of the workflow template workflow = workflow_template.copy() # Update workflow nodes with our parameters # Player image node workflow["391"]["inputs"]["image"] = player_upload["name"] # Pose image node workflow["310"]["inputs"]["image"] = pose_upload["name"] # Shirt image node workflow["636"]["inputs"]["image"] = shirt_upload["name"] # Player name node workflow["471"]["inputs"]["string"] = f"_{player_name}" # Team name node workflow["667"]["inputs"]["string"] = team_name # Age, gender, background prompt node workflow["420"]["inputs"]["string"] = age_gender_bg # Shirt color node workflow["528"]["inputs"]["string"] = f"({shirt_color}:1.2) blank t-shirt, black shorts, " # Style node workflow["422"]["inputs"]["string"] = style # Queue the prompt in ComfyUI prompt_response = queue_prompt(workflow) if "error" in prompt_response: return None, f"Error queuing prompt: {prompt_response['error']}" prompt_id = prompt_response["prompt_id"] # Wait for the processing to complete status = "Generating avatar..." retries = 0 max_retries = 60 # 5 minutes timeout while retries < max_retries: time.sleep(5) progress = check_progress(prompt_id) if "error" in progress: retries += 1 continue if prompt_id in progress and len(progress[prompt_id]["outputs"]) > 0: # Get the output image for node_id, output in progress[prompt_id]["outputs"].items(): if node_id == "308" or node_id == "679": # Save Image nodes image_filename = output.get("images", [{}])[0].get("filename", "") if image_filename: result_image = get_image(image_filename, "", "output") masked_filename = output.get("images", [{}])[0].get("filename", "").replace(".png", "_Masked.png") masked_image = get_image(masked_filename, "", "output") # Return the masked image if available, otherwise the regular image return masked_image if masked_image else result_image, "Avatar generated successfully!" return None, "Completed, but couldn't find output image." retries += 1 status = f"Generating avatar... (attempt {retries}/{max_retries})" return None, "Timed out waiting for the avatar generation to complete." def create_interface(): """Create the Gradio interface for the avatar generator""" with gr.Blocks(title="Football Player Avatar Generator") as demo: gr.Markdown("# Football Player Avatar Generator") gr.Markdown("Create stylized football player avatars from photos") with gr.Row(): with gr.Column(): player_image = gr.Image(label="Upload Player Photo", type="pil") with gr.Row(): pose_image = gr.Image(label="Select Pose Template", type="pil") shirt_image = gr.Image(label="Select Shirt Template", type="pil") player_name = gr.Textbox(label="Player Name", value="Player") team_name = gr.Textbox(label="Team Name", value="Kurjet") age_gender_options = [ "9 year old boy with on light grey studio background, upper body portrait", "10 year old boy with on light grey studio background, upper body portrait", "adult man with on light grey studio background, upper body portrait", "adult woman with on light grey studio background, upper body portrait" ] age_gender_bg = gr.Dropdown(label="Age, Gender & Background", choices=age_gender_options, value=age_gender_options[0]) shirt_color_options = ["black", "red", "blue", "green", "yellow", "white"] shirt_color = gr.Dropdown(label="Shirt Color", choices=shirt_color_options, value="black") style_options = [ "3d pixar character portrait, award winning, 3d animation, octane rendering", "digital painting, detailed, concept art, smooth, sharp focus, illustration, trending on artstation", "cartoon drawing, hand drawn, pencil on paper, sketch art", "watercolor painting, beautiful, smooth, sharp focus, colorful, professional" ] style = gr.Dropdown(label="Art Style", choices=style_options, value=style_options[0]) generate_button = gr.Button("Generate Avatar", variant="primary") with gr.Column(): output_image = gr.Image(label="Generated Avatar") status_text = gr.Textbox(label="Status", interactive=False) # Load default images for pose and shirt try: default_pose = Image.open("/app/ComfyUI/input/pose4.jpg") default_shirt = Image.open("/app/ComfyUI/input/paita2.jpg") pose_image.value = default_pose shirt_image.value = default_shirt except Exception as e: print(f"Error loading default images: {str(e)}") # Set up the button click event generate_button.click( fn=generate_avatar, inputs=[player_image, pose_image, shirt_image, player_name, team_name, age_gender_bg, shirt_color, style], outputs=[output_image, status_text] ) # Add examples if os.path.exists("/app/ComfyUI/input"): example_images = [] for filename in os.listdir("/app/ComfyUI/input"): if filename.lower().endswith(('.png', '.jpg', '.jpeg')): example_images.append(os.path.join("/app/ComfyUI/input", filename)) if example_images: gr.Examples( examples=[[img] for img in example_images[:3]], inputs=[player_image] ) return demo # Check for ComfyUI server and launch the app if __name__ == "__main__": # Check if ComfyUI is running retry_count = 0 max_retries = 5 comfy_running = False while retry_count < max_retries and not comfy_running: try: response = requests.get(f"{COMFY_API}/system_stats") if response.status_code == 200: print("ComfyUI is running, starting Gradio interface...") comfy_running = True break except: pass retry_count += 1 print(f"Waiting for ComfyUI to start... (attempt {retry_count}/{max_retries})") time.sleep(10) if not comfy_running: print("WARNING: Could not connect to ComfyUI server. The application may not function correctly.") # Create and launch the interface demo = create_interface() demo.launch(server_name="0.0.0.0", server_port=7860, share=False)