Update app.py

app.py

@@ -5,6 +5,7 @@ import nltk
 from nltk.tokenize import sent_tokenize
 from PIL import Image
 from io import BytesIO
+import base64
 
 # Ensure NLTK's Punkt tokenizer models are downloaded
 nltk.download('punkt')
@@ -17,51 +18,70 @@ if not HUGGINGFACE_API_KEY:
 
 # Define the Hugging Face Inference API URL for a text-to-image model
 # You can choose a different model if preferred
-API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
+MODEL_NAME = "stabilityai/stable-diffusion-2-1"  # Ensure this model is available for your API key
+API_URL = f"https://api-inference.huggingface.co/models/{MODEL_NAME}"
 
 headers = {
     "Authorization": f"Bearer {HUGGINGFACE_API_KEY}"
 }
 
-def query(payload):
+def query(prompt):
     """
-    Sends a POST request to the Hugging Face Inference API with the given payload.
+    Sends a POST request to the Hugging Face Inference API with the given prompt.
+    Returns the image bytes if successful.
     """
+    payload = {
+        "inputs": prompt,
+        # You can add additional parameters here if needed, such as "parameters": {"num_inference_steps": 50}
+    }
+
     response = requests.post(API_URL, headers=headers, json=payload)
+
     if response.status_code == 200:
-        return response.content  # Image bytes
+        # The response is a list of base64-encoded images
+        response_json = response.json()
+        if isinstance(response_json, list) and len(response_json) > 0:
+            image_data = response_json[0]  # Get the first image
+            # The image data is a data URI: "data:image/png;base64,..."
+            header, encoded = image_data.split(",", 1)
+            image_bytes = base64.b64decode(encoded)
+            return image_bytes
+        else:
+            raise ValueError("Unexpected response format from Hugging Face API.")
     else:
-        raise Exception(f"Request failed with status code {response.status_code}: {response.text}")
+        # Attempt to parse error message from response
+        try:
+            error_info = response.json()
+            error_message = error_info.get('error', 'No error message provided.')
+        except ValueError:
+            error_message = response.text
+        raise Exception(f"Request failed with status code {response.status_code}: {error_message}")
 
 def generate_comic(story):
     """
     Takes a short story, splits it into sentences, generates an image for each sentence,
-    and returns a list of image objects.
+    and returns a list of PIL Image objects.
     """
     # Split the story into sentences
     sentences = sent_tokenize(story)
-    
+
     images = []
     for idx, sentence in enumerate(sentences):
         try:
             print(f"Generating image for sentence {idx+1}/{len(sentences)}: {sentence}")
-            # Prepare the payload for the API
-            payload = {
-                "inputs": sentence,
-                "options": {
-                    "use_cache": False
-                }
-            }
-            # Query the API
-            image_bytes = query(payload)
+            # Query the Hugging Face API
+            image_bytes = query(sentence)
             # Open the image
             image = Image.open(BytesIO(image_bytes)).convert("RGB")
             images.append(image)
         except Exception as e:
             print(f"Error generating image for sentence '{sentence}': {e}")
-            # Optionally, append a placeholder image or skip
-            images.append(None)
-    
+            # Append a placeholder image with error text
+            placeholder = Image.new('RGB', (512, 512), color=(255, 0, 0))
+            draw = Image.Draw.Draw(placeholder)
+            draw.text((10, 10), "Error generating image", fill=(255, 255, 255))
+            images.append(placeholder)
+
     return images
 
 def create_comic_panel(images):
@@ -71,25 +91,25 @@ def create_comic_panel(images):
     """
     # Filter out any None images
     images = [img for img in images if img is not None]
-    
+
     if not images:
         return None
-    
+
     # Define the size for each panel
     panel_width, panel_height = 512, 512  # You can adjust as needed
-    
+
     # Resize images to uniform size
     resized_images = [img.resize((panel_width, panel_height)) for img in images]
-    
+
     # Create a new blank image with enough width to hold all panels
     total_width = panel_width * len(resized_images)
     max_height = panel_height
     comic_image = Image.new('RGB', (total_width, max_height), color=(255, 255, 255))
-    
+
     # Paste each panel image side by side
     for i, img in enumerate(resized_images):
         comic_image.paste(img, (i * panel_width, 0))
-    
+
     return comic_image
 
 def process_story(story):