Upload app.py with huggingface_hub

app.py

@@ -10,8 +10,8 @@ import io
 import base64
 import spaces
 
-# ComfyUI API endpoint
-COMFY_API = "http://127.0.0.1:8188/api"
+# 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"
 
 # Load the workflow template
@@ -23,11 +23,23 @@ except Exception as e:
     print(f"Error loading workflow template: {str(e)}")
     workflow_template = {}
 
+def is_comfyui_running():
+    """Check if ComfyUI server is running and accessible"""
+    try:
+        response = requests.get(f"{COMFY_API}/system_stats", timeout=5)
+        return response.status_code == 200
+    except Exception as e:
+        print(f"ComfyUI server check failed: {str(e)}")
+        return False
+
 def queue_prompt(prompt):
     """Send a prompt to ComfyUI for processing"""
+    if not is_comfyui_running():
+        return {"error": "ComfyUI server is not running"}
+    
     p = {"prompt": prompt}
     try:
-        response = requests.post(f"{COMFY_API}/prompt", json=p)
+        response = requests.post(f"{COMFY_API}/prompt", json=p, timeout=30)
         return response.json()
     except Exception as e:
         print(f"Error queuing prompt: {str(e)}")
@@ -35,8 +47,11 @@ def queue_prompt(prompt):
 
 def get_image(filename, subfolder, folder_type):
     """Get an image from ComfyUI's output folder"""
+    if not is_comfyui_running():
+        return None
+    
     try:
-        response = requests.get(f"{COMFY_API}/view?filename={filename}&subfolder={subfolder}&type={folder_type}")
+        response = requests.get(f"{COMFY_API}/view?filename={filename}&subfolder={subfolder}&type={folder_type}", timeout=30)
         return Image.open(io.BytesIO(response.content))
     except Exception as e:
         print(f"Error getting image {filename}: {str(e)}")
@@ -44,6 +59,9 @@ def get_image(filename, subfolder, folder_type):
 
 def upload_image(image, filename):
     """Upload an image to ComfyUI's input folder"""
+    if not is_comfyui_running():
+        return {"error": "ComfyUI server is not running"}
+    
     try:
         if isinstance(image, str):  # Base64 string
             image_data = base64.b64decode(image.split(",")[1])
@@ -57,16 +75,25 @@ def upload_image(image, filename):
             img_byte_arr.seek(0)
             files = {"image": (filename, img_byte_arr.getvalue())}
         
-        response = requests.post(f"{COMFY_API}/upload/image", files=files)
+        # Use a longer timeout for uploading images
+        response = requests.post(f"{COMFY_API}/upload/image", files=files, timeout=60)
         return response.json()
+    except requests.exceptions.ConnectionError as e:
+        error_msg = f"Connection error when uploading image: {str(e)}"
+        print(error_msg)
+        return {"error": error_msg}
     except Exception as e:
-        print(f"Error uploading image: {str(e)}")
-        return {"error": str(e)}
+        error_msg = f"Error uploading image: {str(e)}"
+        print(error_msg)
+        return {"error": error_msg}
 
 def check_progress(prompt_id):
     """Check the progress of a ComfyUI prompt"""
+    if not is_comfyui_running():
+        return {"error": "ComfyUI server is not running"}
+    
     try:
-        response = requests.get(f"{COMFY_API}/history/{prompt_id}")
+        response = requests.get(f"{COMFY_API}/history/{prompt_id}", timeout=30)
         return response.json()
     except Exception as e:
         print(f"Error checking progress: {str(e)}")
@@ -75,49 +102,70 @@ 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"""
+    if not is_comfyui_running():
+        return None, "Error: ComfyUI server is not running. Please try again later."
+    
     # Upload images to ComfyUI
+    print("Uploading player image...")
     player_upload = upload_image(player_image, "player.png")
+    
+    print("Uploading pose image...")
     pose_upload = upload_image(pose_image, "pose.png")
+    
+    print("Uploading shirt image...")
     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', '')}"
+        error_msg = f"Error uploading images: {player_upload.get('error', '')} {pose_upload.get('error', '')} {shirt_upload.get('error', '')}"
+        print(error_msg)
+        return None, error_msg
+    
+    print("All images uploaded successfully")
     
     # 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
+    try:
+        # 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
+    except KeyError as e:
+        error_msg = f"Error updating workflow parameters: {str(e)}. The workflow structure may have changed."
+        print(error_msg)
+        return None, error_msg
     
     # Queue the prompt in ComfyUI
+    print("Queuing prompt in ComfyUI...")
     prompt_response = queue_prompt(workflow)
     
     if "error" in prompt_response:
-        return None, f"Error queuing prompt: {prompt_response['error']}"
+        error_msg = f"Error queuing prompt: {prompt_response['error']}"
+        print(error_msg)
+        return None, error_msg
     
     prompt_id = prompt_response["prompt_id"]
+    print(f"Prompt queued with ID: {prompt_id}")
     
     # Wait for the processing to complete
     status = "Generating avatar..."
@@ -130,6 +178,7 @@ def generate_avatar(player_image, pose_image, shirt_image, player_name, team_nam
         
         if "error" in progress:
             retries += 1
+            print(f"Error checking progress (retry {retries}/{max_retries}): {progress['error']}")
             continue
             
         if prompt_id in progress and len(progress[prompt_id]["outputs"]) > 0:
@@ -138,17 +187,26 @@ def generate_avatar(player_image, pose_image, shirt_image, player_name, team_nam
                 if node_id == "308" or node_id == "679":  # Save Image nodes
                     image_filename = output.get("images", [{}])[0].get("filename", "")
                     if image_filename:
+                        print(f"Found output image: {image_filename}")
                         result_image = get_image(image_filename, "", "output")
-                        masked_filename = output.get("images", [{}])[0].get("filename", "").replace(".png", "_Masked.png")
+                        masked_filename = image_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!"
+                        if masked_image:
+                            print("Returning masked image")
+                            return masked_image, "Avatar generated successfully!"
+                        elif result_image:
+                            print("Returning regular image")
+                            return result_image, "Avatar generated successfully!"
+                        else:
+                            return None, "Generated image could not be retrieved."
             
             return None, "Completed, but couldn't find output image."
             
         retries += 1
         status = f"Generating avatar... (attempt {retries}/{max_retries})"
+        print(status)
     
     return None, "Timed out waiting for the avatar generation to complete."
 
@@ -193,6 +251,22 @@ def create_interface():
             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:
@@ -229,18 +303,19 @@ def create_interface():
 if __name__ == "__main__":
     # Check if ComfyUI is running
     retry_count = 0
-    max_retries = 5
+    max_retries = 12  # 2 minutes (10 seconds between retries)
     comfy_running = False
     
     while retry_count < max_retries and not comfy_running:
         try:
-            response = requests.get(f"{COMFY_API}/system_stats")
+            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:
-            pass
+        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})")
@@ -248,7 +323,8 @@ if __name__ == "__main__":
     
     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
     demo = create_interface()
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)