Update app.py

app.py

@@ -1,14 +1,20 @@
 import gradio as gr
 import os
 from PIL import Image, ImageDraw
+import re
 import requests
 from io import BytesIO
-import re
+
 from transformers import pipeline as hf_pipeline
 from diffusers import StableDiffusionPipeline
 import torch
 
-story_gen = hf_pipeline("text-generation", model="tiiuae/falcon-7b-instruct")
+# Use a small text2text model (lightweight)
+story_gen = hf_pipeline("text2text-generation", model="google/flan-t5-base")
+
+# 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")
 
 def screenwriter(prompt: str) -> str:
     instructions = f"""
@@ -25,28 +31,12 @@ def screenwriter(prompt: str) -> str:
     - Very short description of the main character's appearance.
     - IMPORTANT!!! ALWAYS use a delimiter '---' to separate the story from the character description.
 
-    Here's an example:
-        User input: "A unicorn named Jeff discovers a mysterious dish."
-        Your output: 
-            "Jeff is wandering around the enchanted forest. 
-            Jeff stumbles upon a glowing dish in the enchanted forest. 
-            Intrigued, Jeff approaches it cautiously, sensing its magical aura. 
-            As Jeff touches the dish, it reveals that it can either dispose of his horns or make them glow. 
-            Jeff is hesitant, but decides to try the dish. 
-            After taking a sip, Jeff notices that his horns started glowing. 
-            The dish disappears and Jeff is alone in the enchanted forest. "
-            ---
-            "Jeff is a majestic unicorn with shimmering white fur and a spiraled golden horn.
-
-    Be creative, compelling, and format the comic book clearly. Do not include your thought process or any additional commentary in the output.
-
     STORY PROMPT: {prompt}
     """
-
     result = story_gen(instructions, max_new_tokens=250)[0]["generated_text"]
     return result
 
-def remove_think_block(text:str):
+def remove_think_block(text: str):
     return re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL).strip()
 
 def parse_screenwriter_output(output: str):
@@ -69,21 +59,14 @@ def error_image(message):
     d.text((10, 250), message, fill=(255, 0, 0))
     return img
 
-
-pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2", torch_dtype=torch.float16)
-pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
-
-
 def illustrator(story: str, character: str):
-
     if not story or not character:
         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 of the scene. No text. Scene: {scene}. Character: {character}'
+        prompt = f"Comic book illustration. Scene: {scene}. Character: {character}. No text, just visual scene."
         try:
             image = pipe(prompt).images[0]
             images.append((image, scene))
@@ -100,20 +83,14 @@ def pipeline(prompt: str):
     return f"{story}\n---\n{character}", images
 
 
-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')
+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')
     generated_story = gr.Button('Generate Story')
-    with gr.Row():
-        story_output = gr.Textbox(label='Screenwriter', lines=5)
-        gallery = gr.Gallery(label='Comic Scenes')
+    story_output = gr.Textbox(label='Generated Story', lines=5)
+    gallery = gr.Gallery(label='Comic Scenes').style(grid=(2, 2))
+
     generated_story.click(pipeline, inputs=story_input, outputs=[story_output, gallery])
 
 if __name__ == "__main__":
-    demo.launch(mcp_server=True)
+    demo.launch()