Update app.py

app.py

@@ -1,58 +1,44 @@
-import gradio as gr
-from random import randint
-from all_models import models  # Import the list of available models
-from externalmod import gr_Interface_load
-import asyncio
-import os
-from threading import RLock
 from flask import Flask, request, jsonify, send_file
 from flask_cors import CORS
+import asyncio
 import tempfile
-
-lock = RLock()
-HF_TOKEN = os.environ.get("HF_TOKEN")
+import os
+from threading import RLock
+from huggingface_hub import InferenceClient
+from all_models import models  # Importing models from all_models
 
 app = Flask(__name__)
 CORS(app)  # Enable CORS for all routes
 
-# Function to load models
-def load_fn(models):
-    global models_load
-    models_load = {}
-    
-    for model in models:
-        if model not in models_load.keys():
-            try:
-                m = gr_Interface_load(f'models/{model}', hf_token=HF_TOKEN)
-            except Exception as error:
-                print(error)
-                m = gr.Interface(lambda: None, ['text'], ['image'])
-            models_load.update({model: m})
+lock = RLock()
+HF_TOKEN = os.environ.get("HF_TOKEN")  # Hugging Face token
 
-load_fn(models)
+inference_timeout = 600  # Set timeout for inference
 
-num_models = 6  # Number of models to load initially
-MAX_SEED = 3999999999
-default_models = models[:num_models]  # Load the first few models for inference
-inference_timeout = 600
+# Function to dynamically load models from the "models" list
+def get_model_from_name(model_name):
+    return model_name if model_name in models else None
 
 # Asynchronous function to perform inference
-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))
+async def infer(client, prompt, seed=1, timeout=inference_timeout, model="prompthero/openjourney-v4"):
+    task = asyncio.create_task(
+        asyncio.to_thread(client.text_to_image, prompt=prompt, seed=seed, model=model)
+    )
     await asyncio.sleep(0)
     try:
         result = await asyncio.wait_for(task, timeout=timeout)
     except (Exception, asyncio.TimeoutError) as e:
         print(e)
-        print(f"Task timed out: {model_str}")
-        if not task.done(): 
+        print(f"Task timed out for model: {model}")
+        if not task.done():
             task.cancel()
         result = None
+    
     if task.done() and result is not None:
         with lock:
             temp_image = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
-            result.save(temp_image.name)  # Save result as a temporary file
+            with open(temp_image.name, "wb") as f:
+                f.write(result)  # Save the result image as a temporary file
         return temp_image.name  # Return the path to the saved image
     return None
 
@@ -62,16 +48,24 @@ def generate_api():
     data = request.get_json()
 
     # Extract required fields from the request
-    model_str = data.get('model_str', default_models[0])  # Default to first model if not provided
     prompt = data.get('prompt', '')
     seed = data.get('seed', 1)
+    model_name = data.get('model', 'prompthero/openjourney-v4')  # Default to "prompthero/openjourney-v4" if not provided
 
     if not prompt:
         return jsonify({"error": "Prompt is required"}), 400
-    
+
+    # Get the model from all_models
+    model = get_model_from_name(model_name)
+    if not model:
+        return jsonify({"error": f"Model '{model_name}' not found in available models"}), 400
+
     try:
+        # Create a generic InferenceClient for the model
+        client = InferenceClient(token=HF_TOKEN)  # Pass Hugging Face token if needed
+
         # Call the async inference function
-        result_path = asyncio.run(infer(model_str, prompt, seed))
+        result_path = asyncio.run(infer(client, prompt, seed, model=model))
         if result_path:
             return send_file(result_path, mimetype='image/png')  # Send back the generated image file
         else: