Update app.py

app.py

@@ -20,12 +20,12 @@ def generate_garment_image(prompt):
     return placeholder_image
 
 def tryon(person_img, garment_prompt, seed, randomize_seed):
-    post_start_time = time.time()
-    if person_img is None or garment_prompt == "":
-        return None, None, "Empty image or prompt"
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    
+    try:
+        post_start_time = time.time()
+        if person_img is None or garment_prompt == "":
+            return None, seed, "Empty image or prompt"
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
     # Generate garment image from prompt
     garment_img = generate_garment_image(garment_prompt)
     
@@ -45,15 +45,76 @@ def start_tryon(person_img, garment_prompt, seed, randomize_seed):
         seed = random.randint(0, MAX_SEED)
     
     # Generate garment image from prompt
-    garment_img = generate_garment_image(garment_prompt)
-    
-    encoded_person_img = cv2.imencode('.jpg', cv2.cvtColor(person_img, cv2.COLOR_RGB2BGR))[1].tobytes()
-    encoded_person_img = base64.b64encode(encoded_person_img).decode('utf-8')
-    encoded_garment_img = cv2.imencode('.jpg', cv2.cvtColor(garment_img, cv2.COLOR_RGB2BGR))[1].tobytes()
-    encoded_garment_img = base64.b64encode(encoded_garment_img).decode('utf-8')
+        # Generate garment image from prompt
+        garment_img = generate_garment_image(garment_prompt)
+        
+        encoded_person_img = cv2.imencode('.jpg', cv2.cvtColor(person_img, cv2.COLOR_RGB2BGR))[1].tobytes()
+        encoded_person_img = base64.b64encode(encoded_person_img).decode('utf-8')
+        encoded_garment_img = cv2.imencode('.jpg', cv2.cvtColor(garment_img, cv2.COLOR_RGB2BGR))[1].tobytes()
+        encoded_garment_img = base64.b64encode(encoded_garment_img).decode('utf-8')
 
-    # Rest of the function remains the same
-    # ...
+        url = "http://" + os.environ['tryon_url'] + "Submit"
+        token = os.environ['token']
+        cookie = os.environ['Cookie']
+        referer = os.environ['referer']
+        headers = {'Content-Type': 'application/json', 'token': token, 'Cookie': cookie, 'referer': referer}
+        data = {
+            "clothImage": encoded_garment_img,
+            "humanImage": encoded_person_img,
+            "seed": seed
+        }
+        
+        response = requests.post(url, headers=headers, data=json.dumps(data), timeout=50)
+        print("post response code", response.status_code)
+        if response.status_code == 200:
+            result = response.json()['result']
+            status = result['status']
+            if status == "success":
+                uuid = result['result']
+                print(uuid)
+        else:
+            return None, seed, f"Error: HTTP {response.status_code}"
+
+        post_end_time = time.time()
+        print(f"post time used: {post_end_time-post_start_time}")
+
+        get_start_time = time.time()
+        time.sleep(9)
+        Max_Retry = 10
+        result_img = None
+        for i in range(Max_Retry):
+            try:
+                url = "http://" + os.environ['tryon_url'] + "Query?taskId=" + uuid
+                response = requests.get(url, headers=headers, timeout=15)
+                print("get response code", response.status_code)
+                if response.status_code == 200:
+                    result = response.json()['result']
+                    status = result['status']
+                    if status == "success":
+                        result = base64.b64decode(result['result'])
+                        result_np = np.frombuffer(result, np.uint8)
+                        result_img = cv2.imdecode(result_np, cv2.IMREAD_UNCHANGED)
+                        result_img = cv2.cvtColor(result_img, cv2.COLOR_RGB2BGR)
+                        info = "Success"
+                        break
+                    elif status == "error":
+                        return None, seed, "Too many users, please try again later"
+                else:
+                    print(response.text)
+                    info = "URL error, please contact the admin"
+            except requests.exceptions.ReadTimeout:
+                print("timeout")
+                info = "Too many users, please try again later"
+            except Exception as err:
+                print(f"Error: {err}")
+            time.sleep(1)
+        get_end_time = time.time()
+        print(f"get time used: {get_end_time-get_start_time}")
+
+        return result_img, seed, info
+    except Exception as e:
+        print(f"An error occurred: {str(e)}")
+        return None, seed, f"An error occurred: {str(e)}"
 
 MAX_SEED = 999999