Update utility/image_generator.py

utility/image_generator.py

@@ -20,25 +20,36 @@ def generate_image_prompts(script):
     
     return prompts
 
-def hf_pipeline(prompt):
-    API_URL = "https://api-inference.huggingface.co/models/Shakker-Labs/AWPortrait-FL"
-    headers = {"Authorization": f"Bearer {os.getenv('HF_API_TOKEN')}"}
-
-    response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
-
-    if response.status_code == 200:
-        return Image.open(io.BytesIO(response.content))  # Return the image directly
-    else:
-        raise Exception(f"Failed to generate image. Status code: {response.status_code}, {response.text}")
+def hf_pipeline(prompt, max_retries=5, delay=30):
+    retries = 0
+    while retries < max_retries:
+        response = requests.post(f"https://api-inference.huggingface.co/models/Shakker-Labs/AWPortrait-FL", 
+                                 json={"inputs": prompt})
+        if response.status_code == 503:
+            print(f"Model is loading, retrying in {delay} seconds...")
+            retries += 1
+            time.sleep(delay)
+        elif response.status_code == 200:
+            return response.json()
+        else:
+            raise Exception(f"Failed to generate image. Status code: {response.status_code}, {response.text}")
+    
+    raise Exception(f"Failed to generate image after {max_retries} retries.")
 
 def generate_images(prompts):
-    image_files = []
-    for idx, prompt in enumerate(prompts):
-        print(f"Generating image for prompt: {prompt}")
-        # Ensure the prompt is processed on the correct device
-        image = hf_pipeline(prompt).images[0]
-        filename = f"generated_image_{idx}.png"
-        image.save(filename)
-        image_files.append(filename)
+    try:
+         image_files = []
+        for idx, prompt in enumerate(prompts):
+            print(f"Generating image for prompt: {prompt}")
+            # Ensure the prompt is processed on the correct device
+            image = hf_pipeline(prompt).images[0]
+            filename = f"generated_image_{idx}.png"
+            image.save(filename)
+            image_files.append(filename)
     
-    return image_files
+        return image_files
+    except Exception as e:
+        print(e)
+        
+        
+