Update app.py

app.py

@@ -13,87 +13,75 @@ from functools import partial
 from tsr.system import TSR
 from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation
 
-#...
+# HF_TOKEN = os.getenv("HF_TOKEN")
 
-css = """
-/*... */
-"""
+if torch.cuda.is_available():
+    device = "cuda:0"
+else:
+    device = "cpu"
 
-with gr.Blocks(theme=CustomTheme(), css=css) as demo:
-    #...
+d = os.environ.get("DEVICE", None)
+if d!= None:
+    device = d 
 
-    with gr.Column():
-        gr.Markdown("**Input Image**")
-        input_image = gr.Image(
-            label="",
-            image_mode="RGBA",
-            sources="upload",
-            type="pil",
-            elem_id="content_image",
-            width=500,  # Увеличена ширина входных изображений
-        )
-        with gr.Column():
-            do_remove_background = gr.Checkbox(
-                label="Remove Background",
-                value=True,
-                elem_id="remove_background",
-            )
-            foreground_ratio = gr.Slider(
-                label="Foreground Ratio",
-                minimum=0.5,
-                maximum=1.0,
-                value=0.85,
-                step=0.05,
-                elem_id="foreground_ratio",
-            )
-            with gr.Row():
-                submit = gr.Button(
-                    label="Generate",
-                    scale=0,
-                    variant="primary",
-                    elem_classes="generate-button",
-                )
-                output_model = gr.Model3D(
-                    label="Output Model",
-                    interactive=False,
-                )
-                output_model2 = gr.Model3D(
-                    label="Output Model",
-                    interactive=False,
-                )
-            with gr.Tab("obj"):
-                output_model = gr.Model3D(
-                    label="Output Model",
-                    interactive=False,
-                )
-            with gr.Tab("glb"):
-                output_model2 = gr.Model3D(
-                    label="Output Model",
-                    interactive=False,
-                )
-            #...
+model = TSR.from_pretrained(
+    "stabilityai/TripoSR",
+    config_name="config.yaml",
+    weight_name="model.ckpt",
+#    token=HF_TOKEN
+)
+model.renderer.set_chunk_size(131072)
+ # Увеличение разрешения входного изображения
+model.to(device)
 
-    with gr.Row():
-        output_model = gr.Model3D(
-            label="Output Model",
-            interactive=False,
-        )
-        output_model2 = gr.Model3D(
-            label="Output Model",
-            interactive=False,
-        )
-   
+rembg_session = rembg.new_session()
 
+def preprocess(input_image, do_remove_background, foreground_ratio):
+    def fill_background(image):
+        image = np.array(image).astype(np.float32) / 255.0
+        image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
+        image = Image.fromarray((image * 255.0).astype(np.uint8))
+        return image
+
+    if do_remove_background:
+        image = image.convert("RGB")
+        image = remove_background(image, rembg_session)
+        image = resize_foreground(image, foreground_ratio)
+        image = fill_background(image)
+    else:
+        if image.mode == "RGBA":
+            image = fill_background(image)
+    return image
+
+def generate(image):
+    scene_codes = model(image, device=device)
+ # Увеличение разрешения входного изображения
+    mesh = model.extract_mesh(scene_codes)[0]
+    mesh = to_gradio_3d_orientation(mesh)
+    mesh_path = tempfile.NamedTemporaryFile(suffix=".obj", delete=False)
+    mesh.export(mesh_path.name)
+ # Увеличение разрешения входного изображения
+    mesh_path2 = tempfile.NamedTemporaryFile(suffix=".glb", delete=False)
+    mesh.export(mesh_path2.name)
+ # Увеличение разрешения входного изображения
+    return mesh_path.name, mesh_path2.name
+
+def run_example(image_pil):
+ # Увеличение разрешения входного изображения
+    preprocessed = preprocess(image_pil, False, 0.9)
+    mesh_name, mesh_name2 = generate(preprocessed) # Увеличение разрешения входного изображения
+    return image_pil, mesh_name, mesh_name2
+
+with gr.Blocks() as demo:
+    gr.Markdown("# **Input Image**")
     with gr.Column():
-        gr.Markdown("**Input Image**")
         input_image = gr.Image(
             label="",
             image_mode="RGBA",
             sources="upload",
             type="pil",
             elem_id="content_image",
-            width=500,  # Увеличена ширина входных изображений
-        )
+        ) # Увеличение разрешения входного изображения
         with gr.Column():
             do_remove_background = gr.Checkbox(
                 label="Remove Background",
@@ -107,113 +95,42 @@ with gr.Blocks(theme=CustomTheme(), css=css) as demo:
                 value=0.85,
                 step=0.05,
                 elem_id="foreground_ratio",
+            ) # Увеличение разрешения входного изображения
+        with gr.Row():
+            submit = gr.Button("Generate")
+            submit.click(
+                fn=check_input_image,
+                inputs=[input_image],
+                outputs=[],
+            ).click(
+                fn=preprocess,
+                inputs=[input_image, do_remove_background, foreground_ratio],
+                outputs=[input_image],
+            ).click(
+                fn=generate,
+                inputs=[input_image],
+                outputs=["mesh_name", "mesh_name2"],
             )
-            with gr.Row():
-                submit = gr.Button(
-                    label="Generate",
-                    scale=0,
-                    variant="primary",
-                    elem_classes="generate-button",
-                )
-                output_model = gr.Model3D(
-                    label="Output Model",
-                    interactive=False,
-                )
-                output_model2 = gr.Model3D(
-                    label="Output Model",
-                    interactive=False,
-                )
-            with gr.Tab("obj"):
-                output_model = gr.Model3D(
-                    label="Output Model",
-                    interactive=False,
-                )
-                output_model2 = gr.Model3D(
-                    label="Output Model",
-                    interactive=False,
-                )
-            with gr.Tab("glb"):
-                output_model2 = gr.Model3D(
-                    label="Output Model",
-                    interactive=False,
-                )
-
-    with gr.Column():
-        gr.Markdown("**Output Models**")
-        output_model = gr.Model3D(label="Output Model", interactive=False)
-        output_model2 = gr.Model3D(label="Output Model", interactive=False)
-
-    with gr.Row():
-        #...
-
-    with gr.Button("Download Example", elem_id="download_example", fn=partial(run_example, inputs=[input_image], outputs=[input_image, output_model]):
-        #...
+        with gr.Tab("obj"):
+            output_model = gr.Model3D(
+                label="Output Model", 
+                interactive=False,
+            )
+        with gr.Tab("glb"):
+            output_model2 = gr.Model3D(
+                label="Output Model", 
+                interactive=False,
+            )
+    #...
 
     with gr.Column():
         gr.Examples(
-            examples=[
-                "examples/1.png",
-                "examples/2.png",
-            ],
+            examples=["examples/1.png"],
             inputs=[input_image],
-            outputs=[output_model, output_model2],
-            fn=partial(run_example),
-            label="Examples",
-            examples_per_page=20,
-        )
-        submit.click(
-            fn=partial(run_example),
-            inputs=[input_image],
-            outputs=[output_model, output_model2],
-        )
-
-    with gr.Button("Generate", elem_id="generate", variant="primary", fn=generate, inputs=[output_model, output_model2]):
-        #...
-
-    with gr.Row():
-        output_model = gr.Model3D(
-            label="Output Model",
-            interactive=False,
-        )
-        output_model2 = gr.Model3D(
-            label="Output Model",
-            interactive=False,
+            outputs=[input_image, output_model, output_model2],
+            fn=run_example,
+            cache_examples=True,
         )
-    )
 
-    with gr.Column():
-        gr.Markdown("**Input Image**")
-        input_image = gr.Image(
-            label="",
-            image_mode="RGBA",
-            sources="upload",
-            type="pil",
-            elem_id="content_image",
-            width=500,  # Увеличена ширина входных изображений
-        )
-        with gr.Column():
-            do_remove_background = gr.Checkbox(
-                label="Remove Background",
-                value=True,
-                elem_id="remove_background",
-            )
-            foreground_ratio = gr.Slider(
-                label="Foreground Ratio",
-                minimum=0.5,
-                maximum=1.0,
-                value=0.85,
-                step=0.05,
-                elem_id="foreground_ratio",
-            )
-            with gr.Row():
-                submit = gr.Button(
-                    label="Generate",
-                    scale=0,
-                    variant="primary",
-                    elem_classes="generate-button",
-                )
-                output_model = gr.Model3D(
-                    label="Output Model",
-                    interactive=False,
-                )
-                output_model2
+demo.queue(max_size=10)
+demo.launch()