app.py

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)