Upload ./app.py with huggingface_hub

app.py

@@ -134,89 +134,59 @@ if BAKE_AVAILEBLE:
 def gen_save_folder(max_size=30):
     os.makedirs('./outputs/app_output', exist_ok=True)
     exists = set(int(_) for _ in os.listdir('./outputs/app_output') if not _.startswith("."))
-    if len(exists) == max_size: 
-        shutil.rmtree(f"./outputs/app_output/0")
-        cur_id = 0
-    else:
-        cur_id = min(set(range(max_size)) - exists)
+    cur_id = min(set(range(max_size)) - exists) if len(exists)<max_size else -1
     if os.path.exists(f"./outputs/app_output/{(cur_id + 1) % max_size}"):
         shutil.rmtree(f"./outputs/app_output/{(cur_id + 1) % max_size}")
-    save_folder = f'./outputs/app_output/{cur_id}'
+        print(f"remove ./outputs/app_output/{(cur_id + 1) % max_size} success !!!")
+    save_folder = f'./outputs/app_output/{max(0, cur_id)}'
     os.makedirs(save_folder, exist_ok=True)
     print(f"mkdir {save_folder} suceess !!!")
     return save_folder
 
-def stage_0_t2i(text, seed, step, save_folder):
-    dst = save_folder + '/img.png'
-    image = worker_t2i(text, seed, step)
-    image.save(dst)
-    img_nobg = worker_xbg(image, force=True)
-    dst = save_folder + '/img_nobg.png'
-    img_nobg.save(dst)
-    return dst
-
-def stage_1_xbg(image, save_folder, force_remove): 
-    if isinstance(image, str):
-        image = Image.open(image)
-    dst =  save_folder + '/img_nobg.png'
-    rgba = worker_xbg(image, force=force_remove)
-    rgba.save(dst)
-    return dst
-
-def stage_2_i2v(image, seed, step, save_folder):
-    if isinstance(image, str):
-        image = Image.open(image)
-    gif_dst = save_folder + '/views.gif'
-    res_img, pils = worker_i2v(image, seed, step)
-    save_gif(pils, gif_dst)
+
+# @spaces.GPU(duration=150)
+def gen_pipe(text, image=None, do_removebg=True, sseed=0, sstep=25, SSEED=0, SSTEP=50, color='face', 
+        bake=False, render=True, max_faces=12000, force=False, front='auto', others=[180], align_times=3):
+    save_folder = gen_save_folder()
+    image_gen = image is not None
+    if not image_gen:
+        image = worker_t2i(text, sseed, sstep)
+        image.save(save_folder + '/img.png')
+    img_nobg = worker_xbg(image, force=do_removebg if image_gen else True)
+    img_nobg.save(save_folder + '/img_nobg.png')
+    yield img_nobg, None, None, None, None, None
+    res_img, pils = worker_i2v(img_nobg, seed=SSEED, steps=SSTEP)
+    save_gif(pils,  save_folder + '/views.gif')
     views_img, cond_img = res_img[0], res_img[1]
     img_array = np.asarray(views_img, dtype=np.uint8)
     show_img = rearrange(img_array, '(n h) (m w) c -> (n m) h w c', n=3, m=2)
-    show_img = show_img[worker_i2v.order, ...]
-    show_img = rearrange(show_img, '(n m) h w c -> (n h) (m w) c', n=2, m=3)
-    show_img = Image.fromarray(show_img) 
-    return views_img, cond_img, show_img
-
-def stage_3_v23(
-    views_pil, 
-    cond_pil, 
-    seed, 
-    save_folder,
-    target_face_count = 30000,
-    texture_color = 'face'
-): 
-    do_texture_mapping = texture_color == 'face'
+    show_img = rearrange(show_img[worker_i2v.order, ...], '(n m) h w c -> (n h) (m w) c', n=2, m=3)
+    show_img = Image.fromarray(show_img)
+    yield img_nobg, show_img, None, None, None, None
+    do_texture_mapping = color == 'face'
     worker_v23(
-        views_pil, 
-        cond_pil, 
-        seed = seed, 
+        views_img, cond_img, seed = SSEED,
         save_folder = save_folder,
-        target_face_count = target_face_count,
+        target_face_count = max_faces,
         do_texture_mapping = do_texture_mapping
     )
-    glb_dst = save_folder + '/mesh.glb' if do_texture_mapping else None
-    obj_dst =  save_folder + '/mesh.obj'
-    obj_dst = save_folder + '/mesh_vertex_colors.obj' # gradio just only can show vertex shading
-    return obj_dst, glb_dst
-
-def stage_3p_baking(save_folder, color, bake, force, front, others, align_times):
-    if color == "face" and bake:
+    glb_v23 = save_folder + '/mesh.glb' if do_texture_mapping else None
+    obj_v23 = save_folder + '/mesh.obj'
+    obj_v23 = save_folder + '/mesh_vertex_colors.obj'
+    yield img_nobg, show_img, obj_v23, glb_v23, None, None
+    glb_dst = None
+    if do_texture_mapping and bake:
         obj_dst = worker_baker(save_folder, force, front, others, align_times)
         glb_dst = obj_dst.replace(".obj", ".glb")
-        return glb_dst
-    else:
-        return None
-
-def stage_4_gif(save_folder, color, bake, render):
-    if not render: return None
-    baked_fld_list = sorted(glob(save_folder + '/view_*/bake/mesh.obj'))
-    obj_dst = baked_fld_list[-1] if len(baked_fld_list)>=1 else save_folder+'/mesh.obj'
-    assert os.path.exists(obj_dst), f"{obj_dst} file not found"
-    gif_dst = obj_dst.replace(".obj", ".gif")
-    worker_gif(obj_dst, gif_dst_path=gif_dst)
-    return gif_dst
-
-
+    yield img_nobg, show_img, obj_v23, glb_v23, glb_dst, None
+    if do_texture_mapping and render: 
+        baked_obj_list = sorted(glob(save_folder + '/view_*/bake/mesh.obj'))
+        obj_dst = baked_obj_list[-1] if len(baked_obj_list)>=1 else save_folder+'/mesh.obj'
+        assert os.path.exists(obj_dst), f"{obj_dst} file not found"
+        gif_dst = obj_dst.replace(".obj", ".gif")
+        worker_gif(obj_dst, gif_dst_path=gif_dst)
+        yield img_nobg, show_img, obj_v23, glb_v23, glb_dst, gif_dst
+        
 def check_image_available(image):
     if image is None:
         return "Please upload image", gr.update()
@@ -252,7 +222,7 @@ def update_mode(mode):
     }[mode]
     face_change = {
         'Vertex color': gr.update(value=120000, maximum=300000),
-        'Moderate': gr.update(value=60000, maximum=300000),
+        'Face color': gr.update(value=60000, maximum=300000),
         'Baking': gr.update(value=10000, maximum=60000)
     }[mode]
     render_change = {
@@ -269,279 +239,133 @@ def update_mode(mode):
 with gr.Blocks() as demo:
     gr.Markdown(CONST_HEADER)
     with gr.Row(variant="panel"):
-        
-        ###### Input region
-        
         with gr.Column(scale=2):
-            
-            ### Text iutput region
-            
             with gr.Tab("Text to 3D"):
                 with gr.Column():
                     text = gr.TextArea('一只黑白相间的熊猫在白色背景上居中坐着，呈现出卡通风格和可爱氛围。', 
                                        lines=3, max_lines=20, label='Input text (within 70 words)')
+                with gr.Row():
+                    gr.Examples(examples=example_ts, inputs=[text], label="Text examples", examples_per_page=10)
+                with gr.Row():
+                    textgen_submit = gr.Button("Generate", variant="primary") 
 
-                    textgen_mode = gr.Radio(
-                        choices=['Vertex color', 'Face color', 'Baking'], 
-                        label="Texture mode",
-                        value='Baking',
-                        interactive=True
-                    )
-                    
-                    with gr.Accordion("Custom settings", open=False):
-                        textgen_color = gr.Radio(choices=["vertex", "face"], label="Color", value="face")
-                        
-                        with gr.Row():
-                            textgen_render = gr.Checkbox(
-                                label="Do Rendering", 
-                                value=True, 
-                                interactive=True
-                            )
-                            textgen_bake = gr.Checkbox(
-                                label="Do Baking", 
-                                value=True if BAKE_AVAILEBLE else False, 
-                                interactive=True if BAKE_AVAILEBLE else False
-                            )
-                            
-                        with gr.Row():
-                            textgen_seed = gr.Number(value=0, label="T2I seed", precision=0, interactive=True)
-                            textgen_SEED = gr.Number(value=0, label="Gen seed", precision=0, interactive=True)
-
-                        textgen_step = gr.Slider(
-                            value=25,
-                            minimum=15,
-                            maximum=50,
-                            step=1,
-                            label="T2I steps",
-                            interactive=True
-                        )
-                        textgen_STEP = gr.Slider(
-                            value=50,
-                            minimum=20,
-                            maximum=80,
-                            step=1,
-                            label="Gen steps",
-                            interactive=True
-                        )
-                        textgen_max_faces =gr.Slider(
-                            value=10000,
-                            minimum=2000,
-                            maximum=60000,
-                            step=1000,
-                            label="Face number limit",
-                            interactive=True
-                        )
-
-                        with gr.Accordion("Baking Options", open=False):
-                            textgen_force_bake = gr.Checkbox(
-                                label="Force (Ignore the degree of matching)", 
-                                value=False, 
-                                interactive=True
-                            )
-                            textgen_front_baking = gr.Radio(
-                                choices=['input image', 'multi-view front view', 'auto'], 
-                                label="Front view baking",
-                                value='auto',
-                                interactive=True,
-                                visible=True
-                            )
-                            textgen_other_views = gr.CheckboxGroup(
-                                choices=['60°', '120°', '180°', '240°', '300°'], 
-                                label="Other views baking",
-                                value=['180°'],
-                                interactive=True,
-                                visible=True
-                            )
-                            textgen_align_times = gr.Slider(
-                                value=3,
-                                minimum=1,
-                                maximum=5,
-                                step=1,
-                                label="Number of alignment attempts per view",
-                                interactive=True
-                            )
-                        
-                    with gr.Row():
-                        textgen_submit = gr.Button("Generate", variant="primary")
-
-                    with gr.Row():
-                        gr.Examples(examples=example_ts, inputs=[text], label="Text examples", examples_per_page=10)
-                        
-                        
-            textgen_mode.change(
-                fn=update_mode,
-                inputs=textgen_mode, 
-                outputs=[textgen_color, textgen_bake, textgen_max_faces, textgen_render]
-            )
-            textgen_color.change(
-                fn=lambda x:[
-                    gr.update(value=(x=='face'), interactive=(x=='face'), visible=(x=='face')),
-                    gr.update(value=(x=='face'), interactive=(x=='face'), visible=(x=='face')),
-                ],
-                inputs=textgen_color, 
-                outputs=[
-                    textgen_bake, 
-                    textgen_render
-                ]
-            )
-            textgen_bake.change(
-                fn= lambda x:[
-                    gr.update(visible=x),
-                    gr.update(visible=x),
-                    gr.update(visible=x),
-                    gr.update(visible=x),
-                    gr.update(value=10000 if x else 120000, minimum=2000, maximum=60000 if x else 300000)
-                ],
-                inputs=textgen_bake, 
-                outputs=[
-                    textgen_front_baking, 
-                    textgen_other_views, 
-                    textgen_align_times, 
-                    textgen_force_bake, 
-                    textgen_max_faces
-                ]
-            )
-            
-            ### Image iutput region
-            
             with gr.Tab("Image to 3D"):
                 with gr.Row():
                     input_image = gr.Image(label="Input image", width=256, height=256, type="pil",
                                            image_mode="RGBA", sources="upload", interactive=True)
                 with gr.Row():
                     alert_message = gr.Markdown("")  # for warning 
-                    
-                imggen_mode = gr.Radio(
-                    choices=['Vertex color', 'Face color', 'Baking'], 
-                    label="Texture mode",
-                    value='Baking',
-                    interactive=True
-                )
-                
-                with gr.Accordion("Custom settings", open=False):
-                    imggen_color = gr.Radio(choices=["vertex", "face"], label="Color", value="face")
-
-                    with gr.Row():
-                        imggen_removebg = gr.Checkbox(
-                            label="Remove Background", 
-                            value=True, 
-                            interactive=True
-                        )
-                        imggen_render = gr.Checkbox(
-                            label="Do Rendering", 
-                            value=True, 
-                            interactive=True
-                        )
-                        imggen_bake = gr.Checkbox(
-                            label="Do Baking", 
-                            value=True if BAKE_AVAILEBLE else False, 
-                            interactive=True if BAKE_AVAILEBLE else False
-                        )
-                    imggen_SEED = gr.Number(value=0, label="Gen seed", precision=0, interactive=True)
-
-                    imggen_STEP = gr.Slider(
-                        value=50,
-                        minimum=20,
-                        maximum=80,
-                        step=1,
-                        label="Gen steps",
+                with gr.Row():
+                    gr.Examples(examples=example_is,  inputs=[input_image], 
+                        label="Img examples", examples_per_page=10)
+                with gr.Row():
+                    removebg = gr.Checkbox(
+                        label="Remove Background", 
+                        value=True, 
                         interactive=True
                     )
-                    imggen_max_faces = gr.Slider(
-                        value=10000,
-                        minimum=2000,
-                        maximum=60000,
-                        step=1000,
-                        label="Face number limit",
+                    imggen_submit = gr.Button("Generate", variant="primary")
+
+            mode = gr.Radio(
+                choices=['Vertex color', 'Face color', 'Baking'], 
+                label="Texture mode",
+                value='Baking',
+                interactive=True
+            )
+
+            with gr.Accordion("Custom settings", open=False):
+                color = gr.Radio(choices=["vertex", "face"], label="Color", value="face")
+                with gr.Row():
+
+                    render = gr.Checkbox(
+                        label="Do Rendering", 
+                        value=True, 
                         interactive=True
                     )
+                    bake = gr.Checkbox(
+                        label="Do Baking", 
+                        value=True if BAKE_AVAILEBLE else False, 
+                        interactive=True if BAKE_AVAILEBLE else False
+                    )
 
-                    with gr.Accordion("Baking Options", open=False):
-                        imggen_force_bake = gr.Checkbox(
-                                label="Force (Ignore the degree of matching)", 
-                                value=False, 
-                                interactive=True
-                            )
-                        imggen_front_baking = gr.Radio(
-                            choices=['input image', 'multi-view front view', 'auto'], 
-                            label="Front view baking",
-                            value='auto',
-                            interactive=True,
-                            visible=True
-                        )
-                        imggen_other_views = gr.CheckboxGroup(
-                            choices=['60°', '120°', '180°', '240°', '300°'], 
-                            label="Other views baking",
-                            value=['180°'],
-                            interactive=True,
-                            visible=True
-                        )
-                        imggen_align_times =gr.Slider(
-                            value=3,
-                            minimum=1,
-                            maximum=5,
-                            step=1,
-                            label="Number of alignment attempts per view",
-                            interactive=True
-                        )
-
-                input_image.change(
-                    fn=check_image_available, 
-                    inputs=input_image, 
-                    outputs=[alert_message, imggen_removebg]
-                )
-                
-                imggen_mode.change(
-                    fn=update_mode,
-                    inputs=imggen_mode, 
-                    outputs=[imggen_color, imggen_bake, imggen_max_faces, imggen_render]
+                with gr.Row():
+                    seed = gr.Number(value=0, label="T2I seed", precision=0, interactive=True)
+                    SEED = gr.Number(value=0, label="Gen seed", precision=0, interactive=True)
+
+                step = gr.Slider(
+                    value=25,
+                    minimum=15,
+                    maximum=50,
+                    step=1,
+                    label="T2I steps",
+                    interactive=True
                 )
-                
-                imggen_color.change(
-                    fn=lambda x:[
-                        gr.update(value=(x=='face'), interactive=(x=='face'), visible=(x=='face')),
-                        gr.update(value=(x=='face'), interactive=(x=='face'), visible=(x=='face'))
-                    ], 
-                    inputs=imggen_color, 
-                    outputs=[
-                        imggen_bake, 
-                        imggen_render
-                    ]
+                STEP = gr.Slider(
+                    value=50,
+                    minimum=20,
+                    maximum=80,
+                    step=1,
+                    label="Gen steps",
+                    interactive=True
                 )
-                
-                imggen_bake.change(
-                    fn= lambda x:[
-                        gr.update(visible=x),
-                        gr.update(visible=x),
-                        gr.update(visible=x),
-                        gr.update(visible=x),
-                        gr.update( value=10000 if x else 120000, minimum=2000, maximum=60000 if x else 300000)
-                    ],
-                    inputs=imggen_bake, 
-                    outputs=[
-                        imggen_front_baking, 
-                        imggen_other_views, 
-                        imggen_align_times,
-                        imggen_force_bake, 
-                        imggen_max_faces
-                    ]
+                max_faces = gr.Slider(
+                    value=10000,
+                    minimum=2000,
+                    maximum=60000,
+                    step=1000,
+                    label="Face number limit",
+                    interactive=True
                 )
 
-                with gr.Row():
-                    imggen_submit = gr.Button("Generate", variant="primary")      
+                with gr.Accordion("Baking Options", open=False):
+                    force_bake = gr.Checkbox(
+                            label="Force (Ignore the degree of matching)", 
+                            value=False, 
+                            interactive=True
+                        )
+                    front_baking = gr.Radio(
+                        choices=['input image', 'multi-view front view', 'auto'], 
+                        label="Front view baking",
+                        value='auto',
+                        interactive=True,
+                        visible=True
+                    )
+                    other_views = gr.CheckboxGroup(
+                        choices=['60°', '120°', '180°', '240°', '300°'], 
+                        label="Other views baking",
+                        value=['180°'],
+                        interactive=True,
+                        visible=True
+                    )
+                    align_times =gr.Slider(
+                        value=1,
+                        minimum=1,
+                        maximum=5,
+                        step=1,
+                        label="Number of alignment attempts per view",
+                        interactive=True
+                    )
+
+            input_image.change(
+                fn=check_image_available, 
+                inputs=input_image, 
+                outputs=[alert_message, removebg]
+            )
+
+            mode.change(
+                fn=update_mode,
+                inputs=mode, 
+                outputs=[color, bake, max_faces, render]
+            )
 
-                with gr.Row():
-                    gr.Examples(examples=example_is,  inputs=[input_image], 
-                        label="Img examples", examples_per_page=10)
-                    
             gr.Markdown(CONST_NOTE)
-                    
+
         ###### Output region
 
         with gr.Column(scale=3):
             with gr.Row():
                 with gr.Column(scale=2):
-                    rem_bg_image = gr.Image(
+                    rembg_image = gr.Image(
                         label="Image without background", 
                         type="pil",
                         image_mode="RGBA", 
@@ -553,7 +377,7 @@ with gr.Blocks() as demo:
                         type="pil", 
                         interactive=False
                     )
-            
+
             result_3dobj = gr.Model3D(
                 clear_color=[0.0, 0.0, 0.0, 0.0],
                 label="OBJ vertex color",
@@ -562,7 +386,7 @@ with gr.Blocks() as demo:
                 camera_position=[90, 90, None],
                 interactive=False
             )
-                
+
             result_3dglb_texture = gr.Model3D(
                 clear_color=[0.0, 0.0, 0.0, 0.0],
                 label="GLB face color",
@@ -578,9 +402,9 @@ with gr.Blocks() as demo:
                 visible=True,
                 camera_position=[90, 90, None],
                 interactive=False)
-            
+
             result_gif = gr.Image(label="GIF", interactive=False)
-                
+
             with gr.Row():
                 gr.Markdown(
                     "Due to Gradio limitations, OBJ files are displayed with vertex shading only, "
@@ -592,79 +416,24 @@ with gr.Blocks() as demo:
     #===============================================================
     # gradio running code
     #===============================================================
-    
-    save_folder = gr.State()
-    cond_image = gr.State()
-    views_image = gr.State()
-    
-    def handle_click(save_folder):
-        if save_folder is None:
-            save_folder = gen_save_folder()
-        return save_folder
+
+    none = gr.State(None)
 
     textgen_submit.click(
-        fn=handle_click,
-        inputs=[save_folder],
-        outputs=[save_folder]
-    ).success(
-        fn=stage_0_t2i, 
-        inputs=[text, textgen_seed, textgen_step, save_folder], 
-        outputs=[rem_bg_image],
-    ).success(
-        fn=stage_2_i2v, 
-        inputs=[rem_bg_image, textgen_SEED, textgen_STEP, save_folder], 
-        outputs=[views_image, cond_image, result_image],
-    ).success(
-        fn=stage_3_v23, 
-        inputs=[views_image, cond_image, textgen_SEED, save_folder, textgen_max_faces, textgen_color], 
-        outputs=[result_3dobj, result_3dglb_texture],
-    ).success(
-        fn=stage_3p_baking, 
-        inputs=[save_folder, textgen_color, textgen_bake,
-               textgen_force_bake, textgen_front_baking, textgen_other_views, textgen_align_times], 
-        outputs=[result_3dglb_baked],
-    ).success(
-        fn=stage_4_gif, 
-        inputs=[save_folder, textgen_color, textgen_bake, textgen_render], 
-        outputs=[result_gif],
-    ).success(lambda: print('Text_to_3D Done ...\n'))
+        fn=gen_pipe,
+        inputs=[text, none, removebg, seed, step, SEED, STEP, color, bake, render, max_faces, force_bake,
+                front_baking, other_views, align_times],
+        outputs=[rembg_image, result_image, result_3dobj, result_3dglb_texture, result_3dglb_baked, result_gif],
+    )
 
-    
     imggen_submit.click(
-        fn=handle_click,
-        inputs=[save_folder],
-        outputs=[save_folder]
-    ).success(
-        fn=stage_1_xbg, 
-        inputs=[input_image, save_folder, imggen_removebg], 
-        outputs=[rem_bg_image],
-    ).success(
-        fn=stage_2_i2v, 
-        inputs=[rem_bg_image, imggen_SEED, imggen_STEP, save_folder], 
-        outputs=[views_image, cond_image, result_image],
-    ).success(
-        fn=stage_3_v23, 
-        inputs=[views_image, cond_image, imggen_SEED, save_folder, imggen_max_faces, imggen_color],
-        outputs=[result_3dobj, result_3dglb_texture],
-    ).success(
-        fn=stage_3p_baking, 
-        inputs=[save_folder, imggen_color, imggen_bake, 
-                imggen_force_bake, imggen_front_baking, imggen_other_views, imggen_align_times], 
-        outputs=[result_3dglb_baked],
-    ).success(
-        fn=stage_4_gif, 
-        inputs=[save_folder, imggen_color, imggen_bake, imggen_render], 
-        outputs=[result_gif],
-    ).success(lambda: print('Image_to_3D Done ...\n'))
-    
-    #===============================================================
-    # start gradio server
-    #===============================================================
-    CONST_PORT = 8080
-    CONST_MAX_QUEUE = 1
-    CONST_SERVER = '0.0.0.0'
+        fn=gen_pipe,
+        inputs=[none, input_image, removebg, seed, step, SEED, STEP, color, bake, render, max_faces, force_bake,
+                front_baking, other_views, align_times],
+        outputs=[rembg_image, result_image, result_3dobj, result_3dglb_texture, result_3dglb_baked, result_gif],
+    )
 
-    demo.queue(max_size=CONST_MAX_QUEUE)
-    demo.launch(server_name=CONST_SERVER, server_port=CONST_PORT)
+    demo.queue(max_size=1)
+    demo.launch(server_name='0.0.0.0', server_port=8080)
     # demo.launch()