Geek7 commited on
Commit
2c8a77f
·
verified ·
1 Parent(s): f47aac7

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -73
app.py CHANGED
@@ -1,77 +1,11 @@
1
- from flask import Flask, request, jsonify, send_file
2
- import gradio as gr
3
- from random import randint
4
- from all_models import models
5
- from externalmod import gr_Interface_load
6
- import asyncio
7
- import os
8
- from threading import RLock
9
- from PIL import Image
10
-
11
- app = Flask(__name__)
12
-
13
- lock = RLock()
14
- HF_TOKEN = os.environ.get("HF_TOKEN")
15
-
16
- # Load models
17
- def load_fn(models):
18
- global models_load
19
- models_load = {}
20
-
21
- for model in models:
22
- if model not in models_load.keys():
23
- try:
24
- m = gr_Interface_load(f'models/{model}', hf_token=HF_TOKEN)
25
- except Exception as error:
26
- print(error)
27
- m = gr.Interface(lambda: None, ['text'], ['image'])
28
- models_load.update({model: m})
29
-
30
- load_fn(models)
31
 
32
- num_models = 6
33
- MAX_SEED = 3999999999
34
- default_models = models[:num_models]
35
- inference_timeout = 600
36
-
37
- # Gradio inference function
38
- async def infer(model_str, prompt, seed=1, timeout=inference_timeout):
39
- kwargs = {"seed": seed}
40
- task = asyncio.create_task(asyncio.to_thread(models_load[model_str].fn, prompt=prompt, **kwargs, token=HF_TOKEN))
41
- await asyncio.sleep(0)
42
- try:
43
- result = await asyncio.wait_for(task, timeout=timeout)
44
- except (Exception, asyncio.TimeoutError) as e:
45
- print(e)
46
- print(f"Task timed out: {model_str}")
47
- if not task.done():
48
- task.cancel()
49
- result = None
50
- if task.done() and result is not None:
51
- with lock:
52
- png_path = "generated_image.png"
53
- result.save(png_path) # Save the result as an image
54
- return png_path
55
- return None
56
-
57
- # API function to perform inference
58
- @app.route('/generate-image', methods=['POST'])
59
- def generate_image():
60
- data = request.get_json()
61
- model_str = data['model_str']
62
- prompt = data['prompt']
63
- seed = data.get('seed', 1)
64
-
65
- # Run Gradio inference
66
- result_path = asyncio.run(infer(model_str, prompt, seed))
67
-
68
- if result_path:
69
- # Send back the generated image file
70
- return send_file(result_path, mimetype='image/png')
71
- else:
72
- return jsonify({"error": "Failed to generate image."}), 500
73
 
 
 
74
 
75
- # Add this block to make sure your app runs when called
76
  if __name__ == "__main__":
77
- app.run(host='0.0.0.0', port=7860) # Run directly
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
 
3
+ import os
4
+ import subprocess
5
 
 
6
  if __name__ == "__main__":
7
+ # Run awake.py in the background
8
+ subprocess.Popen(["python", "wk.py"]) # Start awake.py
9
+
10
+ # Run the Flask app using Gunicorn
11
+ os.system("gunicorn -w 4 -b 0.0.0.0:7860 myapp:myapp") # 4 worker processes