Update app.py

app.py

@@ -38,38 +38,22 @@ def flip_text(prompt, negative_prompt, task, steps, sampler, cfg_scale, seed):
     height = 1024
 
     if task == "Playground v2":
-        playground = str(os.getenv("playground"))
-        with closing(create_connection("wss://ashrafb-arpr.hf.space/queue/join", timeout=60)) as conn:
-            conn.send('{"fn_index":0,"session_hash":""}')
-            conn.send(f'{{"fn_index":0,"data":["{prompt}"],"session_hash":""}}')
-            conn.recv()
-            conn.recv()
-            conn.recv()
-            conn.recv()
-            conn.recv()
-            a = conn.recv()
-            print(">> A:", a)
-            photo = json.loads(a)['output']['data'][0]
-            photo = photo.replace('data:image/jpeg;base64,', '').replace('data:image/png;base64,', '')
-            photo = Image.open(io.BytesIO(base64.decodebytes(bytes(photo, "utf-8"))))
-            return photo
+        ua = UserAgent()
+        headers = {
+            'user-agent': f'{ua.random}'
+        }
+        client = Client("https://ashrafb-arpr.hf.space/", headers=headers)
+        result = client.predict(prompt, fn_index=0)
+        return result
             
     if task == "Artigen v3":
-        artigen = str(os.getenv("artigen"))
-        with closing(create_connection("wss://ashrafb-arv3s.hf.space/queue/join", timeout=60)) as conn:
-            conn.send('{"fn_index":0,"session_hash":""}')
-            conn.send(f'{{"fn_index":0,"data":["{prompt}", 0, "No style"],"session_hash":""}}')
-            conn.recv()
-            conn.recv()
-            conn.recv()
-            conn.recv()
-            conn.recv()
-            a = conn.recv()
-            print(">> A:", a)
-            photo = json.loads(a)['output']['data'][0]
-            photo = photo.replace('data:image/jpeg;base64,', '').replace('data:image/png;base64,', '')
-            photo = Image.open(io.BytesIO(base64.decodebytes(bytes(photo, "utf-8"))))
-            return photo
+        ua = UserAgent()
+        headers = {
+            'user-agent': f'{ua.random}'
+        }
+        client = Client("https://ashrafb-arv3s.hf.space/", headers=headers)
+        result = client.predict("cat",0,"Cinematic", fn_index=0)
+        return result
     try:
         with closing(create_connection("wss://google-sdxl.hf.space/queue/join")) as conn:
             conn.send('{"fn_index":3,"session_hash":""}')