Update app.py

app.py

@@ -13,200 +13,207 @@ 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")
-
-HEADER = """
-"""
-
-if torch.cuda.is_available():
-    device = "cuda:0"
-else:
-    device = "cpu"
-
-d = os.environ.get("DEVICE", None)
-if d!= None:
-    device = d 
-
-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)
-
-rembg_session = rembg.new_session()
-
-
-def check_input_image(input_image):
-    if input_image is None:
-        raise gr.Error("No image uploaded!")
-
-
-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 = input_image.convert("RGB")
-        image = remove_background(image, rembg_session)
-        image = resize_foreground(image, foreground_ratio)
-        image = fill_background(image)
-    else:
-        image = input_image
-        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_path2 = tempfile.NamedTemporaryFile(suffix=".glb", delete=False)
-    mesh.export(mesh_path.name)
-    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, mesn_name2 = generate(preprocessed)
-    return preprocessed, mesh_name, mesn_name2
-
-
-class CustomTheme(gr.themes.Base):
-    def __init__(self):
-        super().__init__()
-        self.primary_hue = "#191a1e"
-        self.background_fill_primary = "#191a1e"
-        self.background_fill_secondary = "#191a1e"
-        self.background_fill_tertiary = "#191a1e"
-        self.text_color_primary = "#FFFFFF"
-        self.text_color_secondary = "#FFFFFF"
-        self.text_color_tertiary = "#FFFFFF"
-        self.input_background_fill = "#191a1e"
-        self.input_text_color = "#FFFFFF"
-
+#...
 
 css = """
-/* Скрываем нижний колонтитул */
-footer {
-    visibility: hidden;
-    height: 0;
-    margin: 0;
-    padding: 0;
-    overflow: hidden;
-}
-@import url('https://fonts.googleapis.com/css2?familyPoppins:wght@400;500;700&display=swap');
-/* Применяем шрифты */
-body, input, button, textarea, select,.gr-button {
-    font-family: 'Poppins', sans-serif;
-    background-color: #191a1e!important;
-    color: #FFFFFF;
-}
-/* Настройки заголовков */
-h1, h2, h3, h4, h5, h6 {
-    font-family: 'Poppins', sans-serif;
-    font-weight: 700;
-    color: #FFFFFF;
-}
-/* Стиль для текстовых полей и кнопок */
-input[type="text"], textarea {
-    background-color: #191a1e!important;
-    color: #FFFFFF;
-    border: 1px solid #FFFFFF;
-}
-/* Цвет кнопки Generate */
-.generate-button {
-    background-color: #5271FF!important;
-    color: #FFFFFF!important;
-    border: none;
-    font-weight: bold;
-}
-.generate-button:hover {
-    background-color: #405BBF!important; /* Цвет при наведении */
-/* Выравнивание элементов */
-.drop-image-container {
-    display: flex;
-    flex-direction: column;
-    align-items: center;
-}
-.drop-image,.processed-image {
-    margin-bottom: 20px;
-}
-.foreground-ratio-container {
-    margin-top: 20px;
-    margin-bottom: 20px;
-}
-.generate-button {
-    margin-top: 20px;
-    margin-left: auto;
-    margin-right: auto;
-}
+/*... */
 """
 
 with gr.Blocks(theme=CustomTheme(), css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown("**Input Image**", elem_classes="prompt-text")
-
-        with gr.Row():
-            input_image = gr.Image(
-                label="",
-                image_mode="RGBA",
-                sources="upload",
-                type="pil",
-                elem_id="content_image",
-                width=500,  # Увеличена ширина входного изображения
+    #...
+
+    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",
             )
-            with gr.Column():
-                do_remove_background = gr.Checkbox(
-                    label="Remove Background",
-                    value=True,
+            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",
                 )
-                foreground_ratio = gr.Slider(
-                    label="Foreground Ratio",
-                    minimum=0.5,
-                    maximum=1.0,
-                    value=0.85,
-                    step=0.05,
+                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,
+                )
+            #...
+
+    with gr.Row():
+        output_model = gr.Model3D(
+            label="Output Model",
+            interactive=False,
+        )
+        output_model2 = gr.Model3D(
+            label="Output Model",
+            interactive=False,
+        )
+    )
 
-        processed_image = gr.Image(label="Processed Image", interactive=False)
-
-        with gr.Row():
-            submit = gr.Button(
-                "Generate",
-                scale=0,
-                variant="primary",
-                elem_classes="generate-button",
-            )
-
-        with gr.Tab("obj"):
-            output_model = gr.Model3D(
-                label="Output Model",
-                interactive=False,
+    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",
             )
-        with gr.Tab("glb"):
-            output_model2 = gr.Model3D(
-                label="Output Model",
-                interactive=False,
+            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,
+                )
+                output_model2 = gr.Model3D(
+                    label="Output Model",
+                    interactive=False,
+                )
+            with gr.Tab("glb"):
+                output_model2 = gr.Model3D(
+                    label="Output Model",
+                    interactive=False,
+                )
 
-    submit.click(fn=check_input_image, inputs=[input_image]).success(
-        fn=preprocess,
-        inputs=[input_image, do_remove_background, foreground_ratio],
-        outputs=[processed_image],
-    ).success(
-        fn=generate,
-        inputs=[processed_image],
-        outputs=[output_model, output_model2],
+    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.Column():
+        gr.Examples(
+            examples=[
+                "examples/1.png",
+                "examples/2.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,
+        )
     )
 
-demo.queue(max_size=10)
-demo.launch() 
+    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