Update app.py

app.py

@@ -1,136 +1,96 @@
 import os
-import requests
 import gradio as gr
+import requests
+from huggingface_hub import InferenceApi
 import nltk
 from nltk.tokenize import sent_tokenize
-from PIL import Image
-from io import BytesIO
-import base64
+from dotenv import load_dotenv
+
+# Load environment variables from .env file if present
+load_dotenv()
 
-# Ensure NLTK's Punkt tokenizer models are downloaded
+# Ensure NLTK sentence tokenizer data is downloaded
 nltk.download('punkt')
 
-# Retrieve Hugging Face API key from environment variables
-HUGGINGFACE_API_KEY = os.getenv('onteddu')
+# Retrieve Hugging Face API key from environment variable
+HF_API_KEY = os.getenv('HUGGINGFACE_API_KEY')
 
-if not HUGGINGFACE_API_KEY:
-    raise ValueError("Please set the HUGGINGFACE_API_KEY environment variable.")
+if not HF_API_KEY:
+    raise ValueError("Please set the 'HUGGINGFACE_API_KEY' environment variable.")
 
-# Define the Hugging Face Inference API URL for a text-to-image model
+# Initialize Hugging Face Inference API with a suitable image generation model
 # You can choose a different model if preferred
-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/black-forest-labs/FLUX.1-schnell"
-
-headers = {
-    "Authorization": f"Bearer {HUGGINGFACE_API_KEY}"
-}
-
-def query(prompt):
-    """
-    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:
-        # 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:
-        # 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}")
+MODEL_NAME = "stabilityai/stable-diffusion-2"
+inference = InferenceApi(repo_id=MODEL_NAME, token=HF_API_KEY)
 
-def generate_comic(story):
+def generate_comic_images(story):
     """
-    Takes a short story, splits it into sentences, generates an image for each sentence,
-    and returns a list of PIL Image objects.
+    Takes a short story as input, splits it into sentences, generates an image for each sentence,
+    and returns the list of image URLs.
     """
     # Split the story into sentences
     sentences = sent_tokenize(story)
-
+    
+    if len(sentences) == 0:
+        return ["No sentences found in the input story."]
+    
     images = []
+    
     for idx, sentence in enumerate(sentences):
+        # Optionally, you can modify the prompt to better suit comic-style images
+        prompt = f"Comic book style illustration of: {sentence}"
+        
         try:
-            print(f"Generating image for sentence {idx+1}/{len(sentences)}: {sentence}")
-            # Query the Hugging Face API
-            image_bytes = query(sentence)
-            # Open the image
-            image = Image.open(BytesIO(image_bytes)).convert("RGB")
-            images.append(image)
+            # Generate image using Hugging Face Inference API
+            output = inference(prompt)
+            
+            # The output can be a URL or binary data depending on the model
+            # Here we assume it's a URL to the generated image
+            if isinstance(output, str):
+                image_url = output
+            elif isinstance(output, dict) and 'image' in output:
+                image_url = output['image']
+            else:
+                image_url = None
+            
+            if image_url:
+                images.append(image_url)
+            else:
+                images.append("Image generation failed for this sentence.")
         except Exception as e:
-            print(f"Error generating image for sentence '{sentence}': {e}")
-            # 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)
-
+            images.append(f"Error: {str(e)}")
+    
     return images
 
 def create_comic_panel(images):
     """
-    Combines a list of PIL Images into a single image arranged horizontally
-    to mimic comic panels.
+    Arranges images in a comic panel layout.
     """
-    # 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
+    # Create a list of Gradio Image components with captions
+    image_components = []
+    for idx, img in enumerate(images):
+        image_components.append(gr.Image(value=img, label=f"Sentence {idx+1}"))
+    
+    return image_components
 
 def process_story(story):
     """
-    Processes the input story and returns a comic panel image.
+    Main processing function to generate images and arrange them in a comic panel.
     """
-    images = generate_comic(story)
-    comic = create_comic_panel(images)
-    return comic
+    images = generate_comic_images(story)
+    image_components = create_comic_panel(images)
+    return image_components
 
-# Define the Gradio interface
+# Define Gradio interface
 iface = gr.Interface(
     fn=process_story,
     inputs=gr.Textbox(lines=10, placeholder="Enter your short story here...", label="Short Story"),
-    outputs=gr.Image(type="pil", label="Generated Comic"),
-    title="Story to Comic Generator",
-    description="Enter a short story, and this app will generate a comic-style image with one panel per sentence.",
+    outputs=gr.Column(),
+    title="Comic Generator",
+    description="Enter a short story, and generate a comic-style image for each sentence using Hugging Face's API.",
     examples=[
-        ["Once upon a time, there was a brave knight. He fought dragons and saved the kingdom."]
-    ],
-    allow_flagging="never"
+        ["Once upon a time, there was a brave knight. He fought dragons and rescued princesses."]
+    ]
 )
 
 if __name__ == "__main__":