Update app.py

app.py

@@ -31,6 +31,7 @@ def flip_text(prompt, negative_prompt, task, steps, sampler, cfg_scale, seed):
     url_sd1 = os.getenv("url_sd1")
     url_sd2 = os.getenv("url_sd2")
     url_sd3 = os.getenv("url_sd3")
+    url_sd4 = os.getenv("url_sd4")
     
     print(task)
     try:
@@ -49,36 +50,50 @@ def flip_text(prompt, negative_prompt, task, steps, sampler, cfg_scale, seed):
             photo = Image.open(io.BytesIO(base64.decodebytes(bytes(photo, "utf-8"))))
             return photo
     except:
-        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
+            print("n_2")
+            with closing(create_connection(f"{url_sd4}", timeout=60)) as conn:
+                conn.send('{"fn_index":0,"session_hash":""}')
+                conn.send(f'{{"data":["{prompt}","[deformed | disfigured], poorly drawn, [bad : wrong] anatomy, [extra | missing | floating | disconnected] limb, (mutated hands and fingers), blurry","dreamshaperXL10_alpha2.safetensors [c8afe2ef]",30,"DPM++ 2M Karras",7,1024,1024,-1],"event_data":null,"fn_index":0,"session_hash":""}}')
+                conn.recv()
+                conn.recv()
+                conn.recv()
+                conn.recv()
+                photo = json.loads(conn.recv())['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
         except:
-            return None
+            print("n_3")
+            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 {