Update app.py

app.py

@@ -16,21 +16,25 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-def remote_text_encoder(prompts):
-    import requests
-    from huggingface_hub import get_token
+def remote_text_encoder(prompts, max_retries=3):
+    from gradio_client import Client
+    import time
     
-    response = requests.post(
-        "https://remote-text-encoder-flux-2.huggingface.co/predict",
-        json={"prompt": prompts},
-        headers={
-            "Authorization": f"Bearer {get_token()}",
-            "Content-Type": "application/json"
-        },
-        timeout=120
-    )
-    prompt_embeds = torch.load(io.BytesIO(response.content))
-    return prompt_embeds
+    for attempt in range(max_retries):
+        try:
+            client = Client("multimodalart/mistral-text-encoder")
+            result = client.predict(
+                prompt=prompts,
+                api_name="/encode_text"
+            )
+            prompt_embeds = torch.load(result[0])
+            return prompt_embeds
+        except Exception as e:
+            print(f"Text encoder attempt {attempt + 1}/{max_retries} failed: {e}")
+            if attempt < max_retries - 1:
+                time.sleep(2)
+            else:
+                raise Exception(f"Text encoder failed after {max_retries} attempts: {e}")
 
 # Load model
 repo_id = "black-forest-labs/FLUX.2-dev"
@@ -80,7 +84,7 @@ def update_dimensions_from_image(image_list):
 def get_duration(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
     num_images = 0 if image_list is None else len(image_list)
     step_duration = 1 + 0.8 * num_images
-    return max(35, num_inference_steps * step_duration + 10)
+    return max(45, num_inference_steps * step_duration + 10)
 
 @spaces.GPU(duration=get_duration)
 def generate_image(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):