Update app.py

app.py

@@ -1,45 +1,59 @@
 from flask import Flask, render_template,request,jsonify,Response  
+import sys
 import os
+import warnings
+import logging
+import urllib3
 import json
+import requests
 from flask_cors import CORS
 import time
 from flask_limiter import Limiter
 from flask_limiter.util import get_remote_address
-import requests
+
+from hf import *
+
+
+
 app = Flask(__name__)
+
 CORS(app)
+import logging
+
+
+
+
+
+total = 0
+pending = 0
 
-def get_client_ip():
-    return request.headers.get(os.getenv('head'), get_remote_address())
 
 limiter = Limiter(
-    key_func=get_client_ip,
-    app=app,
-    default_limits=["6 per minute"]
+    key_func=get_remote_address,
+    default_limits=["8 per minute"]
 )
+limiter.init_app(app)
+
 
 
-users = 0
-userslist = []
 
 @app.route("/")
 def index():
-    global users 
-    global userslist
-    if(request.remote_addr not in userslist):
-        userslist.append(request.remote_addr)
-    users = len(userslist)
-    return render_template('models.html')
+  message = {'Total Requests: ': total, 'Pending Requests: ': pending}
+  return jsonify(message)
+
+
+
+
 
-    
-@app.route("/users")
-def users():
-    return {"total users : " : users} 
 
 @limiter.limit("6 per minute")
 @app.route('/gen', methods=['POST'])
 def Hf():
     
+    global total, pending
+    total += 1
+    pending += 1
     prompt = request.json.get('prompt', '')
     negative = request.json.get('negative', '')
     steps = request.json.get('steps', 20)
@@ -48,28 +62,116 @@ def Hf():
     scale = request.json.get('scale',7)
     model = request.json.get('model','sd3')
     style = request.json.get('style', 'Cinematic')
-    shash = request.json.get('hash', 'Cinematic')
-    def Gen(prompt,negative,steps,width,height,scale,style,model):
-        req = requests.post('https://xyplon-proxy.hf.space/hf/img/gen',headers={
-            'Authorization' : os.getenv('auth')
-        },json={
-            'prompt': prompt,
-            'negative': negative,
-            'steps': steps,
-            'width': width,
-            'height': height,
-            'scale': scale,
-            'model' : model,
-            'style': style, 
-            'hash': shash
-        }, stream=True)
-        if(req.status_code!=200):
-            return "an error occurred! ", 500
-        for chunk in req.iter_lines():
-            yield f'{chunk.decode()}\n'
+    shash = request.json.get('hash')
+    if model == 'rvs4':
+        def generator():
+            global pending
+            while True:  
+                try:
+                    for info in rvx(prompt=prompt, negative=negative, width=width, height=height, scale=scale, steps=steps, style=style, shash=shash):
+                        if "error" in json.dumps(info):
+                            
+                            raise ValueError("Error in response, retrying...")
+                        else:
+                            yield f'data: {json.dumps(info)}\n\n'
+                    pending -= 1
+                    break 
+                except ValueError as e:
+                    continue  
+        return Response(generator(), mimetype="text/event-stream")
+
+    
+
+    elif model == 'mobius':
+        def generator():
+            global pending
+            while True:  
+                try:
+                    for info in mobius(prompt=prompt, negative=negative, width=width, height=height, scale=scale, steps=steps, style=style, shash=shash):
+                        if "error" in json.dumps(info):
+                            raise ValueError("Error in response, retrying...")
+                        else:
+                            yield f'data: {json.dumps(info)}\n\n'
+                    pending -= 1
+                    break 
+                except ValueError as e:
+                    continue  
+        return Response(generator(), mimetype="text/event-stream")
+    
+    elif model == 'sd3':
+        def generator():
+            global pending
+            while True:  
+                try:
+                    for info in sd3(prompt=prompt, negative=negative, width=width, height=height, scale=scale, steps=steps, style=style, shash=shash):
+                        if "error" in json.dumps(info):
+                            raise ValueError("Error in response, retrying...")
+                        else:
+                            yield f'data: {json.dumps(info)}\n\n'
+                    pending -= 1
+                    break 
+                except ValueError as e:
+                    continue  
+        return Response(generator(), mimetype="text/event-stream")
         
-    return Response(Gen(prompt=prompt,negative=negative,steps=steps,width=width,height=height,scale=scale,style=style,model=model), mimetype="text/event-stream")
+    
+    elif model == 'sdflash':
+        def generator():
+            global pending
+            while True:  
+                try:
+                    for info in sdflash(prompt=prompt, negative=negative, width=width, height=height, scale=scale, steps=steps, style=style, shash=shash):
+                        if "error" in json.dumps(info):
+                            raise ValueError("Error in response, retrying...")
+                        else:
+                            yield f'data: {json.dumps(info)}\n\n'
+                    pending -= 1
+                    break 
+                except ValueError as e:
+                    continue  
+        return Response(generator(), mimetype="text/event-stream")
+
+  
+    elif model == 'kivotos':
+        def generator():
+            global pending
+            while True:  
+                try:
+                    for info in kivotos(prompt=prompt, negative=negative, width=width, height=height, scale=scale, steps=steps, style=style, shash=shash):
+                        if "error" in json.dumps(info):
+                            raise ValueError("Error in response, retrying...")
+                        else:
+                            yield f'data: {json.dumps(info)}\n\n'
+                    pending -= 1
+                    break 
+                except ValueError as e:
+                    continue  
+        return Response(generator(), mimetype="text/event-stream")
+        
+    elif model == 'OpenDalle':
+        def generator():
+            global pending
+            while True:  
+                try:
+                    for info in OpenDalle(prompt=prompt, negative=negative, width=width, height=height, scale=scale, steps=steps, style=style, shash=shash):
+                        if "error" in json.dumps(info):
+                            raise ValueError("Error in response, retrying...")
+                        else:
+                            yield f'data: {json.dumps(info)}\n\n'
+                    pending -= 1
+                    break 
+                except ValueError as e:
+                    continue  
+        return Response(generator(), mimetype="text/event-stream")
+        
+    else:
+        pending-=1
+        return jsonify('model not found'), 404
+
+
+
 
 
-if __name__ == "__main__":
-    app.run(debug=True, host='0.0.0.0', port=7860)
+if __name__ == '__main__':
+    #socketio.start_background_task(background_task)
+    app.run(debug=False, host='0.0.0.0', port=7860)