Update app.py

app.py

@@ -1,6 +1,4 @@
 import os
-os.system("pip install git+https://github.com/huggingface/huggingface_hub")
-from huggingface_hub.inference_api import InferenceClient
 import gradio as gr
 import numpy as np
 import random
@@ -8,20 +6,80 @@ import spaces
 import torch
 import tempfile
 import os
+import requests
+import time
+from io import BytesIO
+
 from PIL import Image
 
 token = os.environ["API_TOKEN"]
-model = InferenceClient("black-forest-labs/FLUX.1-schnell", token=token)
+model = "black-forest-labs/FLUX.1-schnell"
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
+def text_to_image(prompt, height, width, seed, num_inference_steps,
+                  model=model,
+                  token=token):
+    """
+    Generate an image from a text prompt using the Hugging Face Inference API.
+    Returns a PIL Image (JPEG) if successful, otherwise returns None after retrying.
+
+    Parameters:
+        prompt (str): The text prompt describing the image.
+        height (int): Height of the generated image.
+        width (int): Width of the generated image.
+        seed (int): Random seed for generation reproducibility.
+        num_inference_steps (int): Number of inference steps.
+        model (str): Hugging Face model identifier.
+        token (str): Hugging Face API token.
+
+    Returns:
+        PIL.Image.Image or None: The generated image as a PIL Image object or None on failure.
+    """
+    api_url = f"https://api-inference.huggingface.co/models/{model}"
+    headers = {
+        "Authorization": f"Bearer {token}"
+    }
+    payload = {
+        "inputs": prompt,
+        "parameters": {
+            "height": height,
+            "width": width,
+            "seed": seed,
+            "num_inference_steps": num_inference_steps
+        }
+    }
+    for attempt in range(1, 4):
+        response = requests.post(api_url, headers=headers, json=payload)
+        if response.status_code == 200:
+            try:
+                image = Image.open(BytesIO(response.content))
+                if image.format != 'JPEG':
+                    with BytesIO() as output:
+                        image.convert("RGB").save(output, format="JPEG")
+                        output.seek(0)
+                        image = Image.open(output)
+                return image
+            except Exception as e:
+                print(f"Error processing the image data: {e}")
+        else:
+            print(f"Attempt {attempt}: Request failed with status code {response.status_code}")
+        
+        if attempt == 1:
+            print("Waiting for 3 seconds before retrying...")
+            time.sleep(3)
+        elif attempt == 2:
+            print("Waiting for 5 seconds before retrying...")
+            time.sleep(5)
+    return None
+
 def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    image = client.text_to_image(
+    image = text_to_image(
         "A girl and a boy dancing in the forest",
         height=height,
         width=width,