Update app.py

app.py

@@ -3,6 +3,7 @@ import requests
 import rembg
 import random
 import gradio as gr
+import numpy
 
 from PIL import Image
 from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
@@ -45,48 +46,56 @@ def inference(input_img, num_inference_steps, guidance_scale, seed ):
     return result
 
 def remove_background(result):
+    print(type(result))
     # Check if the variable is a PIL Image
     if isinstance(result, Image.Image):
+        print('here IF')
         result = rembg.remove(result)
     # Check if the variable is a str filepath
     elif isinstance(result, str):
+        print('here ELIF')
         result = Image.open(result)
         result = rembg.remove(result)
+    elif isinstance(result, numpy.ndarray):
+      print('here ELIF 2')
+      # Convert the NumPy array to a PIL Image
+      result = Image.fromarray(result)
     return result
-    
 
 
+# Create a Gradio interface for the Zero123++ model
 with gr.Blocks() as demo:
-    gr.Markdown("<h1><center> Zero123++ Demo</center></h1>")
-    with gr.Column():
-      input_img = gr.Image(label='Input Image', type='filepath')
-    with gr.Column():
-      output_img = gr.Image(label='Zero123++ Output')
-    with gr.Accordion("Advanced options:", open=False):
-      rm_in_bkg = gr.Checkbox(label='Remove Input Background', )
-      rm_out_bkg = gr.Checkbox(label='Remove Output Background')
-      num_inference_steps = gr.Slider(label="Number of Inference Steps",minimum=15, maximum=100, step=1, value=75, interactive=True)
-      guidance_scale = gr.Slider(label="Classifier Free Guidance Scale",minimum=1.00, maximum=10.00, step=0.1, value=4.0, interactive=True)
-      seed = gr.Number(0, label='Seed')
+    # Display a title
+    gr.HTML("<h1><center> Interactive WebUI : Zero123++ </center></h1>")
+    gr.HTML("<h3><center> A Single Image to Consistent Multi-view Diffusion Base Model</center></h1>")
+    gr.HTML('''<center> <a href='https://arxiv.org/abs/2310.15110' target='_blank'>ArXiv</a> - <a href='https://github.com/SUDO-AI-3D/zero123plus/tree/main' target='_blank'>Code</a> </center>''')
+    with gr.Row():
+      # Input section: Allow users to upload an image
+      with gr.Column():
+          input_img = gr.Image(label='Input Image', type='filepath')
+      
+      # Output section: Display the Zero123++ output image
+      with gr.Column():
+          output_img = gr.Image(label='Zero123++ Output')
+    
+    # Submit button to initiate the inference
     btn = gr.Button('Submit')
 
+    # Advanced options section with accordion for hiding/showing
+    with gr.Accordion("Advanced options:", open=False):
+        rm_in_bkg = gr.Checkbox(label='Remove Input Background')
+        rm_out_bkg = gr.Checkbox(label='Remove Output Background')
+        num_inference_steps = gr.Slider(label="Number of Inference Steps", minimum=15, maximum=100, step=1, value=75, interactive=True)
+        guidance_scale = gr.Slider(label="Classifier Free Guidance Scale", minimum=1.00, maximum=10.00, step=0.1, value=4.0, interactive=True)
+        seed = gr.Number(0, label='Seed')
+    
+
     btn.click(inference, [input_img, num_inference_steps, guidance_scale, seed ], output_img)
     rm_in_bkg.input(remove_background, input_img, input_img)
     rm_out_bkg.input(remove_background, output_img, output_img)
     
-    gr.Examples(
-    examples=[["extinguisher.png", 75, 4.0, 0],
-              ['mushroom.png', 75, 4.0, 0], 
-              ['tianw2.png', 75, 4.0, 0], 
-              ['lysol.png', 75, 4.0, 0], 
-              ['ghost-eating-burger.png', 75, 4.0, 0]
-             ],
-    inputs=[input_img, num_inference_steps, guidance_scale, seed],
-    outputs=output_img,
-    fn=inference,
-    cache_examples=True,
-      )
 
 
-demo.launch()
+demo.launch(debug=True)
+