Update app.py

app.py

@@ -1,20 +1,19 @@
 import gradio as gr
 import os
-from PIL import Image, ImageDraw
 import re
+from PIL import Image, ImageDraw
 import requests
 from io import BytesIO
+from huggingface_hub import InferenceClient
+from dotenv import load_dotenv
 
-from transformers import pipeline as hf_pipeline
-from diffusers import StableDiffusionPipeline
-import torch
+load_dotenv()
 
-# Use a small text2text model (lightweight)
-story_gen = hf_pipeline("text2text-generation", model="google/flan-t5-base")
+hf_token = os.getenv("HF_API_TOKEN")
+if not hf_token:
+    raise ValueError("Set your HF_API_TOKEN environment variable before running.")
 
-# Use a lighter Stable Diffusion model for images
-pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
-pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+client = InferenceClient(token=hf_token)
 
 def screenwriter(prompt: str) -> str:
     instructions = f"""
@@ -33,10 +32,15 @@ def screenwriter(prompt: str) -> str:
 
     STORY PROMPT: {prompt}
     """
-    result = story_gen(instructions, max_new_tokens=250)[0]["generated_text"]
-    return result
 
-def remove_think_block(text: str):
+    response = client.text_generation(
+        model="tiiuae/falcon-7b-instruct", 
+        inputs=instructions, 
+        max_new_tokens=250
+    )
+    return response[0]['generated_text']
+
+def remove_think_block(text:str):
     return re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL).strip()
 
 def parse_screenwriter_output(output: str):
@@ -64,11 +68,17 @@ def illustrator(story: str, character: str):
         raise ValueError('Could not parse story or character from input.')
 
     scenes = [s.strip() for s in story.split('.') if s.strip()]
+
     images = []
     for idx, scene in enumerate(scenes):
-        prompt = f"Comic book illustration. Scene: {scene}. Character: {character}. No text, just visual scene."
+        prompt = f'Comic book illustration of the scene. No text. Scene: {scene}. Character: {character}'
         try:
-            image = pipe(prompt).images[0]
+            response = client.text_to_image(
+                model="stabilityai/stable-diffusion-2", 
+                inputs=prompt
+            )
+            image_url = response['generated_image_url']
+            image = Image.open(BytesIO(requests.get(image_url).content))
             images.append((image, scene))
         except Exception as e:
             images.append((error_image(f'Error: {str(e)}'), f'Error in scene {idx + 1}'))
@@ -82,15 +92,20 @@ def pipeline(prompt: str):
     images = illustrator(story, character)
     return f"{story}\n---\n{character}", images
 
-
-with gr.Blocks(title='Comic Generator') as demo:
-    gr.Markdown("# Comic Generator\nGenerate a comic book from your idea!")
-    story_input = gr.Textbox(label='Story Prompt', placeholder='A unicorn named Jeff discovers a mysterious dish')
+with gr.Blocks(theme=gr.themes.Ocean(),
+               title='Comic Generator') as demo:
+    gr.Markdown(
+        '''
+        # Comic Generator
+        Generates a comic off of your prompt.
+        ''')
+    with gr.Row():
+        story_input = gr.Textbox(label='Story Prompt', placeholder='A unicorn named Jeff discovers a mysterious dish')
     generated_story = gr.Button('Generate Story')
-    story_output = gr.Textbox(label='Generated Story', lines=5)
-    gallery = gr.Gallery(label='Comic Scenes').style(grid=(2, 2))
-
+    with gr.Row():
+        story_output = gr.Textbox(label='Screenwriter', lines=5)
+        gallery = gr.Gallery(label='Comic Scenes')
     generated_story.click(pipeline, inputs=story_input, outputs=[story_output, gallery])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(mcp_server=True)