Update utility/image_generator.py

utility/image_generator.py

@@ -27,12 +27,23 @@ 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 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 = pipeline(prompt).images[0]
+        image = hf_pipeline(prompt).images[0]
         filename = f"generated_image_{idx}.png"
         image.save(filename)
         image_files.append(filename)