Update app.py

app.py

@@ -7,6 +7,9 @@ from websocket import create_connection
 from deep_translator import GoogleTranslator
 from langdetect import detect
 import os
+from PIL import Image
+import io
+import base64
 
 
 def flip_text(prompt, negative_prompt, task, steps, sampler, cfg_scale, seed):
@@ -27,38 +30,56 @@ def flip_text(prompt, negative_prompt, task, steps, sampler, cfg_scale, seed):
     height = 1024
     url_sd1 = os.getenv("url_sd1")
     url_sd2 = os.getenv("url_sd2")
+    url_sd3 = os.getenv("url_sd3")
     
     print(task)
- 
-    if task == 'Stable Diffusion XL 1.0':
-        model = 'sd_xl_base_1.0'
-    if task == 'Crystal Clear XL':
-        model = '[3d] crystalClearXL_ccxl_97637'
-    if task == 'Juggernaut XL':
-        model = '[photorealistic] juggernautXL_version2_113240'
-    if task == 'DreamShaper XL':
-        model = '[base model] dreamshaperXL09Alpha_alpha2Xl10_91562'
-    if task == 'SDXL Niji':
-        model = '[midjourney] sdxlNijiV51_sdxlNijiV51_112807'
-    if task == 'Cinemax SDXL':
-        model = '[movie] cinemaxAlphaSDXLCinema_alpha1_107473'
-    if task == 'NightVision XL':
-        model = '[photorealistic] nightvisionXLPhotorealisticPortrait_beta0702Bakedvae_113098'
-    
-    negative = negative_prompt
-    
     try:
-        with closing(create_connection(f"{url_sd1}")) as conn:
-            conn.send('{"fn_index":231,"session_hash":""}')
-            conn.send(f'{{"data":["task()","{prompt}","{negative}",[],{steps},"{sampler}",false,false,1,1,{cfg},{seed},-1,0,0,0,false,{width},{height},false,0.7,2,"Lanczos",0,0,0,"Use same sampler","","",[],"None",true,"{model}","Automatic",null,null,null,false,false,"positive","comma",0,false,false,"","Seed","",[],"Nothing","",[],"Nothing","",[],true,false,false,false,0,null,null,false,null,null,false,null,null,false,50,[],"","",""],"event_data":null,"fn_index":231,"session_hash":""}}')
-            print(conn.recv())
-            print(conn.recv())
-            print(conn.recv())
-            print(conn.recv())
-            photo = f"{url_sd2}" + str(json.loads(conn.recv())['output']['data'][0][0]["name"])
-        return photo
+        print('n_1')
+        with closing(create_connection(f"{url_sd3}", timeout=60)) as conn:
+            conn.send('{"fn_index":3,"session_hash":""}')
+            conn.send(f'{{"data":["{text}, 4k photo","[deformed | disfigured], poorly drawn, [bad : wrong] anatomy, [extra | missing | floating | disconnected] limb, (mutated hands and fingers), blurry",7.5,"(No style)"],"event_data":null,"fn_index":3,"session_hash":""}}')
+            while True:
+                status = json.loads(conn.recv())['msg']
+                print(status)
+                if status == 'estimation':
+                    continue
+                if status == 'process_starts':
+                    break
+            photo = json.loads(conn.recv())['output']['data'][0][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
     except:
-        return None
+        print(n_2)
+        if task == 'Stable Diffusion XL 1.0':
+            model = 'sd_xl_base_1.0'
+        if task == 'Crystal Clear XL':
+            model = '[3d] crystalClearXL_ccxl_97637'
+        if task == 'Juggernaut XL':
+            model = '[photorealistic] juggernautXL_version2_113240'
+        if task == 'DreamShaper XL':
+            model = '[base model] dreamshaperXL09Alpha_alpha2Xl10_91562'
+        if task == 'SDXL Niji':
+            model = '[midjourney] sdxlNijiV51_sdxlNijiV51_112807'
+        if task == 'Cinemax SDXL':
+            model = '[movie] cinemaxAlphaSDXLCinema_alpha1_107473'
+        if task == 'NightVision XL':
+            model = '[photorealistic] nightvisionXLPhotorealisticPortrait_beta0702Bakedvae_113098'
+        
+        negative = negative_prompt
+        
+        try:
+            with closing(create_connection(f"{url_sd1}")) as conn:
+                conn.send('{"fn_index":231,"session_hash":""}')
+                conn.send(f'{{"data":["task()","{prompt}","{negative}",[],{steps},"{sampler}",false,false,1,1,{cfg},{seed},-1,0,0,0,false,{width},{height},false,0.7,2,"Lanczos",0,0,0,"Use same sampler","","",[],"None",true,"{model}","Automatic",null,null,null,false,false,"positive","comma",0,false,false,"","Seed","",[],"Nothing","",[],"Nothing","",[],true,false,false,false,0,null,null,false,null,null,false,null,null,false,50,[],"","",""],"event_data":null,"fn_index":231,"session_hash":""}}')
+                print(conn.recv())
+                print(conn.recv())
+                print(conn.recv())
+                print(conn.recv())
+                photo = f"{url_sd2}" + str(json.loads(conn.recv())['output']['data'][0][0]["name"])
+            return photo
+        except:
+            return None
 
 css = """
 #generate {