ui fix

app.py

@@ -47,6 +47,7 @@ def serve_static(filename):
 
 async def process_fal_request(request_id, model_endpoint, fal_arguments, api_key):
     """Process FAL API request asynchronously"""
+    handler = None
     try:
         print(f"[DEBUG] Starting async processing for request {request_id}")
         print(f"[DEBUG] Model endpoint: {model_endpoint}")
@@ -72,24 +73,47 @@ async def process_fal_request(request_id, model_endpoint, fal_arguments, api_key
             'result': None
         }
         
-        # Collect logs asynchronously
+        # Collect logs asynchronously with better error handling
         print(f"[DEBUG] Starting to collect events...")
         event_count = 0
-        async for event in handler.iter_events(with_logs=True):
-            event_count += 1
-            print(f"[DEBUG] Event #{event_count}: {type(event).__name__}")
-            
-            if hasattr(event, 'logs') and event.logs:
-                for log in event.logs:
-                    log_msg = log.get("message", "")
-                    print(f"[DEBUG] Log: {log_msg}")
-                    active_requests[request_id]['logs'].append(log_msg)
+        max_events = 1000  # Prevent infinite loops
+        
+        try:
+            async for event in handler.iter_events(with_logs=True):
+                event_count += 1
+                print(f"[DEBUG] Event #{event_count}: {type(event).__name__}")
+                
+                if hasattr(event, 'logs') and event.logs:
+                    for log in event.logs:
+                        log_msg = log.get("message", "")
+                        print(f"[DEBUG] Log: {log_msg}")
+                        active_requests[request_id]['logs'].append(log_msg)
+                
+                # Safety check to prevent infinite loops
+                if event_count >= max_events:
+                    print(f"[WARNING] Max events ({max_events}) reached, breaking loop")
+                    break
+        except asyncio.CancelledError:
+            print(f"[WARNING] Event iteration cancelled for request {request_id}")
+            raise
+        except RuntimeError as e:
+            if "Event loop is closed" in str(e):
+                print(f"[WARNING] Event loop closed during iteration, attempting to continue...")
+            else:
+                print(f"[WARNING] Runtime error during event iteration: {e}")
+        except Exception as iter_error:
+            print(f"[WARNING] Error during event iteration: {iter_error}")
+            # Continue to try to get the result even if event iteration fails
         
         print(f"[DEBUG] Total events collected: {event_count}")
         
-        # Get the final result
+        # Get the final result with timeout to prevent hanging
         print(f"[DEBUG] Getting final result...")
-        result = await handler.get()
+        try:
+            result = await asyncio.wait_for(handler.get(), timeout=60)  # 60 second timeout
+        except asyncio.TimeoutError:
+            print(f"[ERROR] Timeout waiting for result")
+            raise Exception("Timeout waiting for FAL API result")
         
         print(f"[DEBUG] Result received: {json.dumps(result, indent=2)[:500] if result else 'None'}...")
         
@@ -105,6 +129,11 @@ async def process_fal_request(request_id, model_endpoint, fal_arguments, api_key
         print(f"[DEBUG] Async processing completed for request {request_id}")
         return result
         
+    except asyncio.CancelledError:
+        print(f"[WARNING] Task cancelled for request {request_id}")
+        active_requests[request_id]['status'] = 'error'
+        active_requests[request_id]['error'] = 'Task was cancelled'
+        raise
     except Exception as e:
         print(f"[ERROR] Exception in process_fal_request: {str(e)}")
         import traceback
@@ -117,17 +146,23 @@ def run_async_task(request_id, model_endpoint, fal_arguments, api_key):
     """Run async task with proper event loop management"""
     print(f"[DEBUG run_async_task] Starting for request {request_id}")
     
-    # Always create a fresh event loop for each task to avoid conflicts
+    # Create and run event loop in a way that prevents premature closure
     loop = None
     try:
         # Create a new event loop for this thread
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
         
+        # Enable debug mode to get better error messages
+        loop.set_debug(False)  # Set to True for more verbose debugging
+        
         # Run the async task in the new loop
         result = loop.run_until_complete(process_fal_request(request_id, model_endpoint, fal_arguments, api_key))
         print(f"[DEBUG run_async_task] Completed with result keys: {result.keys() if result else 'None'}")
         
+        # Ensure all async generators are closed properly
+        loop.run_until_complete(asyncio.sleep(0))  # Process any remaining callbacks
+        
     except Exception as e:
         print(f"[ERROR run_async_task] Failed: {str(e)}")
         import traceback
@@ -137,22 +172,48 @@ def run_async_task(request_id, model_endpoint, fal_arguments, api_key):
             active_requests[request_id]['status'] = 'error'
             active_requests[request_id]['error'] = str(e)
     finally:
-        # Always close and cleanup the event loop
+        # More careful cleanup to avoid closing loop while still in use
         if loop:
             try:
-                # Cancel any remaining tasks
-                pending = asyncio.all_tasks(loop)
-                for task in pending:
-                    task.cancel()
-                # Run until all tasks are cancelled
-                if pending:
-                    loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
+                # Give time for any final async operations to complete
+                if not loop.is_closed():
+                    # Run a small delay to let any final callbacks execute
+                    loop.run_until_complete(asyncio.sleep(0.5))
+                    
+                    # Get all tasks that are still pending
+                    pending = asyncio.all_tasks(loop)
+                    
+                    # Cancel them gracefully
+                    for task in pending:
+                        task.cancel()
+                    
+                    # Wait for all tasks to be cancelled, with a timeout
+                    if pending:
+                        loop.run_until_complete(
+                            asyncio.wait_for(
+                                asyncio.gather(*pending, return_exceptions=True),
+                                timeout=5.0
+                            )
+                        )
+            except asyncio.TimeoutError:
+                print(f"[WARNING] Timeout during task cleanup")
             except Exception as cleanup_error:
                 print(f"[WARNING] Error during loop cleanup: {cleanup_error}")
             finally:
-                loop.close()
-                # Clear the event loop from the thread
-                asyncio.set_event_loop(None)
+                # Close the loop only after everything is done
+                try:
+                    if not loop.is_closed():
+                        # Stop the loop first to ensure it's not running
+                        loop.stop()
+                        # Give it a moment to stop
+                        loop.run_until_complete(asyncio.sleep(0))
+                        # Now close it
+                        loop.close()
+                except Exception as close_error:
+                    print(f"[WARNING] Error closing loop: {close_error}")
+                finally:
+                    # Clear the event loop from the thread
+                    asyncio.set_event_loop(None)
 
 @app.route('/api/generate', methods=['POST'])
 def generate():