Update app.py

app.py

@@ -1,19 +1,21 @@
 import gradio as gr
 import os
-import re
 from PIL import Image, ImageDraw
-import requests
+import re
 from io import BytesIO
+
 from huggingface_hub import InferenceClient
-from dotenv import load_dotenv
+from diffusers import StableDiffusionPipeline
+import torch
 
-load_dotenv()
+client = InferenceClient()
 
-hf_token = os.getenv("HF_API_TOKEN")
-if not hf_token:
-    raise ValueError("Set your HF_API_TOKEN environment variable before running.")
+pipe = StableDiffusionPipeline.from_pretrained(
+    "CompVis/stable-diffusion-v1-4",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+)
+pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
-client = InferenceClient(token=hf_token)
 
 def screenwriter(prompt: str) -> str:
     instructions = f"""
@@ -34,15 +36,18 @@ def screenwriter(prompt: str) -> str:
     """
 
     response = client.text_generation(
-        model="tiiuae/falcon-7b-instruct", 
-        inputs=instructions, 
-        max_new_tokens=250
+        model="tiiuae/falcon-rw-1b",  # free-tier supported model
+        prompt=instructions,
+        max_new_tokens=250,
+        temperature=0.7
     )
-    return response[0]['generated_text']
+    return response
 
-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):
     cleaned_output = remove_think_block(output)
     delimiter = '---'
@@ -57,34 +62,32 @@ def parse_screenwriter_output(output: str):
         character = lines[-1]
         return story, character
 
+
 def error_image(message):
     img = Image.new("RGB", (512, 512), color=(255, 255, 255))
     d = ImageDraw.Draw(img)
     d.text((10, 250), message, fill=(255, 0, 0))
     return img
 
+
 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 style illustration. No text. Scene: {scene}. Character: {character}"
         try:
-            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))
+            image = pipe(prompt).images[0]
             images.append((image, scene))
         except Exception as e:
             images.append((error_image(f'Error: {str(e)}'), f'Error in scene {idx + 1}'))
     return images
 
-def pipeline(prompt: str):
+
+def comic_pipeline(prompt: str):
     output = screenwriter(prompt)
     story, character = parse_screenwriter_output(output)
     if not story or not character:
@@ -92,20 +95,17 @@ def pipeline(prompt: str):
     images = illustrator(story, character)
     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.Blocks(theme=gr.themes.Ocean(), title='Comic Generator') as demo:
+    gr.Markdown("# Comic Generator\nGive a prompt and get a comic!")
     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')
+    generate_btn = gr.Button('Generate Comic')
     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])
+        story_output = gr.Textbox(label='Screenwriter Output', lines=6)
+        gallery = gr.Gallery(label='Comic Scenes').style(grid=[2], height='auto')
+    generate_btn.click(comic_pipeline, inputs=story_input, outputs=[story_output, gallery])
+
 
 if __name__ == "__main__":
-    demo.launch(mcp_server=True)
+    demo.launch()