Update app.py

app.py

@@ -2,28 +2,26 @@ import gradio as gr
 from random import randint
 from all_models import models
 from datetime import datetime
-import logging
-
-# Set up logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 def get_current_time():
     now = datetime.now()
-    current_time = now.strftime("%Y-%m-%d %H:%M:%S")
-    return current_time
+    now2 = now
+    current_time = now2.strftime("%Y-%m-%d %H:%M:%S")
+    ki = f'{current_time}'
+    return ki
 
-def load_model(model_str):
-    if model_str not in models_load:
-        try:
-            m = gr.load(f'models/{model_str}')
-            logging.info(f"Model {model_str} loaded successfully.")
-        except Exception as error:
-            logging.error(f"Failed to load model {model_str}: {error}")
-            m = gr.Interface(lambda txt: None, ['text'], ['image'])
-        models_load[model_str] = m
-    return models_load[model_str]
+def load_fn(models):
+    global models_load
+    models_load = {}
+    for model in models:
+        if model not in models_load.keys():
+            try:
+                m = gr.load(f'models/{model}')
+            except Exception as error:
+                m = gr.Interface(lambda txt: None, ['text'], ['image'])
+            models_load.update({model: m})
 
-models_load = {}
+load_fn(models)
 
 num_models = len(models)
 default_models = models[:num_models]
@@ -43,21 +41,13 @@ def gen_fn(model_str, prompt, max_retries=10):
     while retries < max_retries:
         try:
             noise = str(randint(0, 9999999))
-            model = load_model(model_str)
-            result = model(f'{prompt} {noise}')
-            logging.info(f"Image generated successfully with model {model_str} and prompt '{prompt}'.")
+            result = models_load[model_str](f'{prompt} {noise}')
             return result
-        except TimeoutError as te:
-            logging.error(f"Timeout error with model {model_str} and prompt '{prompt}': {te}")
-        except ConnectionError as ce:
-            logging.error(f"Connection error with model {model_str} and prompt '{prompt}': {ce}")
         except Exception as e:
-            logging.error(f"General error with model {model_str} and prompt '{prompt}': {e}")
-        
-        retries += 1
-        if retries >= max_retries:
-            logging.error(f"Failed to generate image after {max_retries} retries with model {model_str} and prompt '{prompt}'.")
-            raise Exception(f"Failed to generate image after {max_retries} retries.")
+            print(f"Error generating image: {e}")
+            retries += 1
+            if retries >= max_retries:
+                raise Exception(f"Failed to generate image after {max_retries} retries.")
     
     return None
 
@@ -247,5 +237,5 @@ body {
 with gr.Blocks(css=custom_css) as demo: 
     make_me()
 
-demo.queue(concurrency_count=40)  # Adjust this value based on your resources
+demo.queue(concurrency_count=500)
 demo.launch()