Update app.py

app.py

@@ -36,7 +36,7 @@ MAX_SEED = 3999999999
 default_models = models[:num_models]
 inference_timeout = 600
 
-# Inference function to generate image
+# Inference function with generate_api embedded
 async def infer(model_str, prompt, seed=1, timeout=inference_timeout):
     kwargs = {"seed": seed}
     task = asyncio.create_task(asyncio.to_thread(models_load[model_str].fn, prompt=prompt, **kwargs, token=HF_TOKEN))
@@ -56,14 +56,7 @@ async def infer(model_str, prompt, seed=1, timeout=inference_timeout):
         return png_path
     return None
 
-# Generate API function that calls the async infer function
-def generate_api(model_str, prompt, seed=1):
-    result = asyncio.run(infer(model_str, prompt, seed))
-    if result:
-        return result  # Path to the generated image
-    return None
-
-# Flask route to handle predictions
+# Flask API to call the generate_api function
 @app.route('/predict', methods=['POST'])
 def predict():
     data = request.get_json()
@@ -76,9 +69,9 @@ def predict():
     prompt = data['prompt']
     seed = data.get('seed', 1)
 
-    # Generate the image using the model
+    # Make the asynchronous call to the infer function within the Flask route
     try:
-        image_path = generate_api(model_str, prompt, seed)
+        image_path = asyncio.run(infer(model_str, prompt, seed))  # Directly call infer function here
         if image_path:
             return send_file(image_path, mimetype='image/png')
         else:
@@ -91,5 +84,5 @@ if __name__ == '__main__':
     app.run(debug=True)
 
 # You can optionally launch the Gradio interface in parallel
-iface = gr.Interface(fn=generate_api, inputs=["text", "text", "number"], outputs="file")
+iface = gr.Interface(fn=infer, inputs=["text", "text", "number"], outputs="file")
 iface.launch(show_api=True, share=True)