Update app.py

app.py

@@ -1,4 +1,4 @@
-import spaces
+# import spaces
 import os
 import datetime
 import einops
@@ -22,6 +22,28 @@ from myutils.misc import load_dreambooth_lora, rand_name
 from myutils.wavelet_color_fix import wavelet_color_fix
 from annotator.retinaface import RetinaFaceDetection
 
+from io import BytesIO
+import base64
+import re
+
+# Regex pattern to match data URI scheme
+data_uri_pattern = re.compile(r'data:image/(png|jpeg|jpg|webp);base64,')
+
+def readb64(b64):
+    # Remove any data URI scheme prefix with regex
+    b64 = data_uri_pattern.sub("", b64)
+    # Decode and open the image with PIL
+    img = Image.open(BytesIO(base64.b64decode(b64)))
+    return img
+    
+# convert from PIL to base64
+def writeb64(image):
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    b64image = base64.b64encode(buffered.getvalue())
+    b64image_str = b64image.decode("utf-8")
+    return b64image_str
+
 use_pasd_light = False
 face_detector = RetinaFaceDetection()
 
@@ -84,7 +106,7 @@ def resize_image(image_path, target_height):
         #resized_img.save(output_path)
         return resized_img
 
-@spaces.GPU(enable_queue=True)
+# @spaces.GPU(enable_queue=True)
 def inference(input_image, prompt, a_prompt, n_prompt, denoise_steps, upscale, alpha, cfg, seed):
     input_image = resize_image(input_image, 512)
     process_size = 768
@@ -138,55 +160,18 @@ def inference(input_image, prompt, a_prompt, n_prompt, denoise_steps, upscale, a
             print(e)
             image = Image.new(mode="RGB", size=(512, 512))
     
-    # Convert and save the image as JPEG
-    image.save(f'result_{timestamp}.jpg', 'JPEG')
-
-    # Convert and save the image as JPEG
-    input_image.save(f'input_{timestamp}.jpg', 'JPEG')
-    
-    return (f"input_{timestamp}.jpg", f"result_{timestamp}.jpg"), f"result_{timestamp}.jpg"
+    return writeb64(image)
 
 title = "Pixel-Aware Stable Diffusion for Real-ISR"
 description = "Gradio Demo for PASD Real-ISR. To use it, simply upload your image, or click one of the examples to load them."
 article = "<a href='https://github.com/yangxy/PASD' target='_blank'>Github Repo Pytorch</a>"
-#examples=[['samples/27d38eeb2dbbe7c9.png'],['samples/629e4da70703193b.png']]
-
-css = """
-#col-container{
-    margin: 0 auto;
-    max-width: 720px;
-}
-#project-links{
-    margin: 0 0 12px !important;
-    column-gap: 8px;
-    display: flex;
-    justify-content: center;
-    flex-wrap: nowrap;
-    flex-direction: row;
-    align-items: center;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.HTML(f"""
-        <h2 style="text-align: center;">
-            PASD Magnify
-        </h2>
-        <p style="text-align: center;">
-            Pixel-Aware Stable Diffusion for Realistic Image Super-resolution and Personalized Stylization
-        </p>
-        <p id="project-links" align="center">
-            <a href='https://github.com/yangxy/PASD'><img src='https://img.shields.io/badge/Project-Page-Green'></a> <a href='https://huggingface.co/papers/2308.14469'><img src='https://img.shields.io/badge/Paper-Arxiv-red'></a>
-        </p> 
-        <p style="margin:12px auto;display: flex;justify-content: center;">
-            <a href="https://huggingface.co/spaces/fffiloni/PASD?duplicate=true"><img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-lg.svg" alt="Duplicate this Space"></a>
-        </p>
-        
-        """)
+
+
+with gr.Blocks() as demo:
+    with gr.Column():
         with gr.Row():
             with gr.Column():
-                input_image = gr.Image(type="filepath", sources=["upload"], value="samples/frog.png")
+                input_image = gr.Textbox()
                 prompt_in = gr.Textbox(label="Prompt", value="Frog")
                 with gr.Accordion(label="Advanced settings", open=False):
                     added_prompt = gr.Textbox(label="Added Prompt", value='clean, high-resolution, 8k, best quality, masterpiece')
@@ -198,8 +183,7 @@ with gr.Blocks(css=css) as demo:
                     seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
                 submit_btn = gr.Button("Submit")
             with gr.Column():
-                b_a_slider = ImageSlider(label="B/A result", position=0.5)
-                file_output = gr.File(label="Downloadable image result")
+                output_image = gr.Textbox()
                 
     submit_btn.click(
         fn = inference,
@@ -210,9 +194,6 @@ with gr.Blocks(css=css) as demo:
             upsample_scale, condition_scale,
             classifier_free_guidance, seed
         ],
-        outputs = [
-            b_a_slider,
-            file_output
-        ]
+        outputs = output_image
     )
 demo.queue().launch()