Update app.py

app.py

@@ -9,20 +9,35 @@ import logging
 from datetime import datetime
 
 # Configure logging
-logging.basicConfig(level=logging.DEBUG)
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 class Config:
     RUNPOD_API_KEY = os.getenv("RUNPOD_API_KEY", "")
-    RUNPOD_ENDPOINT = os.getenv("RUNPOD_ENDPOINT", "").replace("/run", "/runsync")  # Force sync endpoint
+    RUNPOD_ENDPOINT = os.getenv("RUNPOD_ENDPOINT", "").rstrip('/')  # Clean endpoint URL
     PORT = 7860
-    TIMEOUT = 120
+    TIMEOUT = 90
     MAX_RETRIES = 2
     IMAGE_SIZE = (768, 768)
 
+    @classmethod
+    def validate_endpoint(cls):
+        """Ensure endpoint has correct format"""
+        if not cls.RUNPOD_ENDPOINT:
+            raise ValueError("RUNPOD_ENDPOINT not set")
+        
+        # Normalize endpoint URL
+        if not cls.RUNPOD_ENDPOINT.endswith(('/run', '/runsync')):
+            cls.RUNPOD_ENDPOINT = f"{cls.RUNPOD_ENDPOINT}/runsync"
+        
+        if not cls.RUNPOD_ENDPOINT.startswith('https://api.runpod.ai/v2/'):
+            raise ValueError("Invalid endpoint format. Should be: https://api.runpod.ai/v2/ENDPOINT_ID/[run|runsync]")
+
 class RunPodClient:
     async def generate_image(self, prompt, negative_prompt=""):
-        """Enhanced response handler with direct status checking"""
+        """Handle both sync and async endpoints"""
+        Config.validate_endpoint()
+        
         for attempt in range(Config.MAX_RETRIES):
             try:
                 async with aiohttp.ClientSession() as session:
@@ -37,58 +52,63 @@ class RunPodClient:
                         }
                     }
 
+                    logger.debug(f"Sending request to: {Config.RUNPOD_ENDPOINT}")
+                    
                     async with session.post(
                         Config.RUNPOD_ENDPOINT,
                         headers={"Authorization": f"Bearer {Config.RUNPOD_API_KEY}"},
                         json=payload,
                         timeout=Config.TIMEOUT
                     ) as response:
+                        
+                        # Handle text/plain responses
+                        content_type = response.headers.get('Content-Type', '')
+                        if 'text/plain' in content_type:
+                            text_response = await response.text()
+                            logger.error(f"Text response received: {text_response}")
+                            raise ValueError(f"Unexpected text response: {text_response[:100]}...")
+                        
                         data = await response.json()
-                        logger.debug(f"Full API response: {data}")
+                        logger.debug(f"API response: {data}")
 
-                        # Handle job ID response (async)
-                        if data.get("id"):
+                        # Handle async job response
+                        if data.get("id") and not Config.RUNPOD_ENDPOINT.endswith('/runsync'):
                             job_id = data["id"]
-                            status_url = f"{Config.RUNPOD_ENDPOINT.replace('/runsync', '')}/status/{job_id}"
+                            status_url = f"{Config.RUNPOD_ENDPOINT.rsplit('/', 1)[0]}/status/{job_id}"
                             return await self._check_job_status(session, status_url)
                         
-                        # Handle direct output (sync)
-                        if data.get("status") == "COMPLETED":
-                            if isinstance(data.get("output"), list):
-                                return data["output"][0], None
-                            return None, "❌ Completed but no output URL found"
-                        
-                        return None, f"❌ Unexpected response: {data.get('error', str(data))}"
+                        # Handle direct output
+                        if data.get("status") == "COMPLETED" and data.get("output"):
+                            return data["output"][0], None
+                            
+                        return None, f"❌ Unexpected response format: {data}"
 
             except Exception as e:
                 logger.error(f"Attempt {attempt + 1} failed: {str(e)}")
                 if attempt == Config.MAX_RETRIES - 1:
                     return None, f"🔥 Connection error: {str(e)}"
-                await asyncio.sleep(2 ** attempt)
+                await asyncio.sleep(1)
 
     async def _check_job_status(self, session, status_url):
         """Poll job status for async endpoints"""
-        retries = 10
-        while retries > 0:
+        for _ in range(10):  # Max 10 checks
             async with session.get(
                 status_url,
                 headers={"Authorization": f"Bearer {Config.RUNPOD_API_KEY}"},
                 timeout=Config.TIMEOUT
             ) as response:
                 data = await response.json()
-                logger.debug(f"Job status check: {data}")
-
+                
                 if data.get("status") == "COMPLETED":
-                    if isinstance(data.get("output"), list) and len(data["output"]) > 0:
+                    if data.get("output"):
                         return data["output"][0], None
-                    return None, "❌ Job completed but no output URL"
+                    return None, "❌ Job completed but no output"
                 
                 if data.get("status") == "FAILED":
                     return None, f"❌ Job failed: {data.get('error', 'Unknown error')}"
                 
                 await asyncio.sleep(2)
-                retries -= 1
-
+        
         return None, "⌛ Timeout waiting for job completion"
 
 def create_interface():
@@ -118,6 +138,7 @@ def create_interface():
                 return None, "❌ Please enter a prompt"
             
             try:
+                Config.validate_endpoint()
                 client = RunPodClient()
                 image_url, error = await client.generate_image(prompt, negative)
                 
@@ -142,11 +163,11 @@ def create_interface():
 
 if __name__ == "__main__":
     try:
+        # Validate configuration before starting
         if not Config.RUNPOD_API_KEY:
-            raise ValueError("RUNPOD_API_KEY not set")
-        if not Config.RUNPOD_ENDPOINT:
-            raise ValueError("RUNPOD_ENDPOINT not set")
-            
+            raise ValueError("RUNPOD_API_KEY environment variable not set")
+        Config.validate_endpoint()
+        
         app = create_interface()
         app.launch(
             server_name="0.0.0.0",