Upload app.py with huggingface_hub

app.py

@@ -12,7 +12,13 @@ import spaces
 
 # ComfyUI API endpoint - use localhost instead of 127.0.0.1
 COMFY_API = "http://localhost:8188/api"
-WORKFLOW_PATH = "/app/workflows/Workflow_12_11.json"
+
+# Get workflow path from environment variable or use default
+WORKFLOW_PATH = os.environ.get("WORKFLOW_PATH", "/tmp/workflows/Workflow_12_11.json")
+
+# Input and output directories (configured in ComfyUI via start.sh)
+INPUT_DIR = "/tmp/comfyui_input"
+OUTPUT_DIR = "/tmp/comfyui_output"
 
 # Load the workflow template
 try:
@@ -26,12 +32,21 @@ except Exception as e:
 def is_comfyui_running():
     """Check if ComfyUI server is running and accessible"""
     try:
-        response = requests.get(f"{COMFY_API}/system_stats", timeout=5)
+        response = requests.get(f"{COMFY_API}/system_stats", timeout=2)
         return response.status_code == 200
     except Exception as e:
         print(f"ComfyUI server check failed: {str(e)}")
         return False
 
+def wait_for_comfyui(timeout=60):
+    """Wait for ComfyUI to be available, with a timeout"""
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        if is_comfyui_running():
+            return True
+        time.sleep(2)
+    return False
+
 def queue_prompt(prompt):
     """Send a prompt to ComfyUI for processing"""
     if not is_comfyui_running():
@@ -39,7 +54,7 @@ def queue_prompt(prompt):
     
     p = {"prompt": prompt}
     try:
-        response = requests.post(f"{COMFY_API}/prompt", json=p, timeout=30)
+        response = requests.post(f"{COMFY_API}/prompt", json=p, timeout=10)
         return response.json()
     except Exception as e:
         print(f"Error queuing prompt: {str(e)}")
@@ -51,7 +66,7 @@ def get_image(filename, subfolder, folder_type):
         return None
     
     try:
-        response = requests.get(f"{COMFY_API}/view?filename={filename}&subfolder={subfolder}&type={folder_type}", timeout=30)
+        response = requests.get(f"{COMFY_API}/view?filename={filename}&subfolder={subfolder}&type={folder_type}", timeout=10)
         return Image.open(io.BytesIO(response.content))
     except Exception as e:
         print(f"Error getting image {filename}: {str(e)}")
@@ -70,13 +85,18 @@ def upload_image(image, filename):
             if isinstance(image, np.ndarray):
                 image = Image.fromarray(image)
             
+            # Resize image if it's too large
+            max_size = (1024, 1024)
+            if image.width > max_size[0] or image.height > max_size[1]:
+                image.thumbnail(max_size, Image.LANCZOS)
+            
             img_byte_arr = io.BytesIO()
-            image.save(img_byte_arr, format='PNG')
+            image.save(img_byte_arr, format='PNG', optimize=True)
             img_byte_arr.seek(0)
             files = {"image": (filename, img_byte_arr.getvalue())}
         
-        # Use a longer timeout for uploading images
-        response = requests.post(f"{COMFY_API}/upload/image", files=files, timeout=60)
+        # Use a shorter timeout for uploading images
+        response = requests.post(f"{COMFY_API}/upload/image", files=files, timeout=20)
         return response.json()
     except requests.exceptions.ConnectionError as e:
         error_msg = f"Connection error when uploading image: {str(e)}"
@@ -93,7 +113,7 @@ def check_progress(prompt_id):
         return {"error": "ComfyUI server is not running"}
     
     try:
-        response = requests.get(f"{COMFY_API}/history/{prompt_id}", timeout=30)
+        response = requests.get(f"{COMFY_API}/history/{prompt_id}", timeout=10)
         return response.json()
     except Exception as e:
         print(f"Error checking progress: {str(e)}")
@@ -102,8 +122,10 @@ def check_progress(prompt_id):
 @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"""
+    # First check if ComfyUI is running
     if not is_comfyui_running():
-        return None, "Error: ComfyUI server is not running. Please try again later."
+        if not wait_for_comfyui(timeout=30):
+            return None, "Error: ComfyUI server is not available. Please try again in a few minutes."
     
     # Upload images to ComfyUI
     print("Uploading player image...")
@@ -170,7 +192,7 @@ def generate_avatar(player_image, pose_image, shirt_image, player_name, team_nam
     # Wait for the processing to complete
     status = "Generating avatar..."
     retries = 0
-    max_retries = 60  # 5 minutes timeout
+    max_retries = 30  # Reduce timeout to 2.5 minutes
     
     while retries < max_retries:
         time.sleep(5)
@@ -216,6 +238,18 @@ def create_interface():
         gr.Markdown("# Football Player Avatar Generator")
         gr.Markdown("Create stylized football player avatars from photos")
         
+        # Add ComfyUI status indicator at the top
+        comfy_status = gr.Textbox(label="ComfyUI Server Status", interactive=False, value="Checking...")
+        
+        def check_comfy_server():
+            if is_comfyui_running():
+                return "✅ ComfyUI server is running"
+            else:
+                return "❌ ComfyUI server is not running - the app may not work correctly"
+        
+        # Auto-check status when loading
+        demo.load(fn=check_comfy_server, outputs=comfy_status)
+        
         with gr.Row():
             with gr.Column():
                 player_image = gr.Image(label="Upload Player Photo", type="pil")
@@ -246,32 +280,20 @@ def create_interface():
                 ]
                 style = gr.Dropdown(label="Art Style", choices=style_options, value=style_options[0])
                 
+                # Add manual server check button
+                check_comfy_button = gr.Button("Check ComfyUI Server Status")
+                check_comfy_button.click(fn=check_comfy_server, outputs=comfy_status)
+                
                 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)
-                
-                # Add ComfyUI status indicator
-                comfy_status = gr.Textbox(label="ComfyUI Server Status", interactive=False)
-                check_comfy_button = gr.Button("Check ComfyUI Server Status")
-                
-                def check_comfy_server():
-                    if is_comfyui_running():
-                        return "✅ ComfyUI server is running"
-                    else:
-                        return "❌ ComfyUI server is not running"
-                
-                check_comfy_button.click(
-                    fn=check_comfy_server,
-                    inputs=[],
-                    outputs=[comfy_status]
-                )
         
         # 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")
+            default_pose = Image.open(os.path.join(INPUT_DIR, "pose4.jpg"))
+            default_shirt = Image.open(os.path.join(INPUT_DIR, "paita2.jpg"))
             pose_image.value = default_pose
             shirt_image.value = default_shirt
         except Exception as e:
@@ -285,11 +307,11 @@ def create_interface():
         )
         
         # Add examples
-        if os.path.exists("/app/ComfyUI/input"):
+        if os.path.exists(INPUT_DIR):
             example_images = []
-            for filename in os.listdir("/app/ComfyUI/input"):
+            for filename in os.listdir(INPUT_DIR):
                 if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
-                    example_images.append(os.path.join("/app/ComfyUI/input", filename))
+                    example_images.append(os.path.join(INPUT_DIR, filename))
             
             if example_images:
                 gr.Examples(
@@ -299,32 +321,10 @@ def create_interface():
     
     return demo
 
-# Check for ComfyUI server and launch the app
+# Start the app
 if __name__ == "__main__":
-    # Check if ComfyUI is running
-    retry_count = 0
-    max_retries = 12  # 2 minutes (10 seconds between retries)
-    comfy_running = False
-    
-    while retry_count < max_retries and not comfy_running:
-        try:
-            print(f"Checking if ComfyUI is running (attempt {retry_count+1}/{max_retries})...")
-            response = requests.get(f"{COMFY_API}/system_stats", timeout=5)
-            if response.status_code == 200:
-                print("ComfyUI is running, starting Gradio interface...")
-                comfy_running = True
-                break
-        except Exception as e:
-            print(f"ComfyUI check failed: {str(e)}")
-        
-        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.")
-        print("Starting Gradio interface anyway...")
-    
-    # Create and launch the interface
+    # Create and launch the interface without waiting for ComfyUI
+    # The interface will handle ComfyUI availability checks when needed
+    print("Starting Gradio interface...")
     demo = create_interface()
     demo.launch(server_name="0.0.0.0", server_port=7860, share=False)