Update app.py

app.py

@@ -1,12 +1,13 @@
 import gradio as gr
 import requests
 import os
-from PIL import Image
+from PIL import Image, ImageOps
 from io import BytesIO
 from tqdm import tqdm
 import time
 import numpy as np
 import potrace
+import base64
 
 # Defining the repository information and the trigger word
 repo = "artificialguybr/LineAniRedmond-LinearMangaSDXL-V2"
@@ -27,7 +28,7 @@ constraints = (
 def generate_image(prompt):
     print("Generating image with prompt:", prompt)
     api_url = f"https://api-inference.huggingface.co/models/{repo}"
-    token = os.getenv("API_TOKEN")
+    token = os.getenv("HF_TOKEN")
     headers = {
         "Authorization": f"Bearer {token}"
     }
@@ -54,19 +55,19 @@ def generate_image(prompt):
         if response.status_code == 200:
             print("Image generation successful!")
             # Open the image and convert it to grayscale
+            image = Image.open(BytesIO(response.content)).convert("L")
             # Convert the grayscale image to a binary image
             image = image.point(lambda p: 255 if p > 128 else 0, mode='1')
+            # Invert the image for potrace (black on white)
+            image = ImageOps.invert(image)
             # Convert the PIL image to a numpy array
             array = np.array(image)
-            # Invert the array for potrace (black on white)
-            array = 255 - array
             # Trace the bitmap to SVG
             bitmap = potrace.Bitmap(array)
             path = bitmap.trace()
             # Generate SVG data
             svg_data = path.to_svg()
             return svg_data
-    
         elif response.status_code == 503:
             time.sleep(1)
             pbar.update(1)
@@ -77,23 +78,31 @@ def generate_image(prompt):
             print("API Error:", response.status_code)
             raise Exception(f"API Error: {response.status_code}")
 
-iface = gr.Interface(
-    fn=generate_image,
-    inputs=gr.Textbox(lines=2, placeholder="Describe the subject here..."),
-    outputs="text",  # SVG data as text
-    title="LineArt XL Image Generator",
-    description=(
-        "Powered by the generous GPU time from Redmond.AI, this LORA, fine-tuned on SD XL 1.0, excels at creating Lineart-themed images across a wide range of subjects. "
-        "Optimized for 1024x1024 resolution, it incorporates the specific tag 'lineart,LineAniAF' directly in the HF Space for ease of use. "
-        "If you appreciate this model and wish to support, consider a donation via [Patreon](https://www.patreon.com/user?u=81570187) or [Ko-fi](https://ko-fi.com/artificialguybr). "
-        "Stay updated on new models by following on [Twitter](https://twitter.com/artificialguybr)."
-    ),
-    examples=[
-        ["a soaring eagle"], 
-        ["a running cheetah"], 
-        ["a majestic oak tree"]
-    ]
-)
+with gr.Blocks() as demo:
+    gr.Markdown("# LineArt XL Image Generator")
+    prompt = gr.Textbox(lines=2, placeholder="Describe the subject here...")
+    generate_button = gr.Button("Generate Image")
+    
+    with gr.Row():
+        svg_display = gr.HTML()
+        download_button = gr.Download()
+
+    def display_svg(svg_data):
+        # Encode SVG data in base64 for data URI
+        svg_bytes = svg_data.encode('utf-8')
+        svg_b64 = base64.b64encode(svg_bytes).decode('utf-8')
+        # Create HTML content with an iframe
+        html_content = f'<iframe src="data:image/svg+xml;base64,{svg_b64}" width="500" height="500"></iframe>'
+        return html_content, svg_bytes
+
+    generate_button.click(
+        generate_image, 
+        inputs=prompt, 
+        outputs=None
+    ).then(
+        display_svg, 
+        inputs=generate_image,
+        outputs=[svg_display, download_button]
+    )
 
-print("Launching Gradio interface...")
-iface.launch()
+demo.launch()