Update app.py

app.py

@@ -2,6 +2,7 @@ import torch
 import requests
 from PIL import Image
 from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
+import rembg
 
 # Load the pipeline
 pipeline = DiffusionPipeline.from_pretrained(
@@ -18,17 +19,59 @@ pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
 )
 pipeline.to('cuda:0')
 
-# Download an example image.
-cond = Image.open(requests.get("https://d.skis.ltd/nrp/sample-data/lysol.png", stream=True).raw)
 
-# Run the pipeline!
-result = pipeline(cond, num_inference_steps=75).images[0]
-# for general real and synthetic images of general objects
-# usually it is enough to have around 28 inference steps
-# for images with delicate details like faces (real or anime)
-# you may need 75-100 steps for the details to construct
+def inference(input_img, num_inference_steps, guidance_scale, seed ):
+    # Download an example image.
+    cond = Image.open(input_img)
+    
+    # Run the pipeline!
+    #result = pipeline(cond, num_inference_steps=75).images[0]
+    result = pipeline(cond, num_inference_steps=num_inference_steps, 
+                  guidance_scale=guidance_scale, 
+                  generator=torch.Generator(pipeline.device).manual_seed(seed)).images[0]
 
-result.show()
-result.save("output.png")
+    # for general real and synthetic images of general objects
+    # usually it is enough to have around 28 inference steps
+    # for images with delicate details like faces (real or anime)
+    # you may need 75-100 steps for the details to construct
+    
+    #result.show()
+    #result.save("output.png")
+    return result
 
+def remove_background(result):
+    result = rembg.remove(result)
+    return result
+
+
+import gradio as gr
+
+with gr.Blocks() as demo:
+    gr.Markdown("<h1><center> Zero123++ Demo</center></h1>")
+    with gr.Column():
+      input_img = gr.Image(label='Input Image', tyoe='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')
+    btn = gr.Button('Submit')
+
+    btn.click(inference, [input_img, num_inference_steps, guidance_scale, seed ], output_img)
+    rm_in_bkg.input(remove_background, input_img, output_img)
+    rm_out_bkg.input(remove_background, output_img, output_img)
+    
+    gr.Examples(
+    examples=[["one.jpg"],['two.jpg'], ['three.jpg']],
+    inputs=input_img,
+    outputs=output_img,
+    fn=dummy,
+    cache_examples=True,
+      )
+
+
+demo.launch()