Update app.py

app.py

@@ -13,124 +13,194 @@ 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")
-
-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)
+# DEVICE
+DEVICE_ENV = os.environ.get("DEVICE", None)
+DEVICE = DEVICE_ENV if DEVICE_ENV else ("cuda:0" if torch.cuda.is_available() else "cpu")
 
+# MODEL
+MODEL_NAME = "stabilityai/TripoSR"
+CONFIG_NAME = "config.yaml"
+WEIGHT_NAME = "model.ckpt"
+# HF_TOKEN = os.getenv("HF_TOKEN")  # Если нужен токен
+
+# REMBG
 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
+# THEME
+HEADER = """
+"""
+
+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?family=Poppins: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;
+}
+"""
+
+def check_input_image(input_image):
+    if input_image is None:
+        raise gr.Error("No image uploaded!")
 
+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
+
+def preprocess(input_image, do_remove_background, foreground_ratio):
     if do_remove_background:
-        image = image.convert("RGB")
+        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)
- # Увеличение разрешения входного изображения
+    model = TSR.from_pretrained(
+        MODEL_NAME,
+        config_name=CONFIG_NAME,
+        weight_name=WEIGHT_NAME,
+        # token=HF_TOKEN
+    )
+    model.renderer.set_chunk_size(131072)
+    model.to(DEVICE)
+    
+    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_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, mesh_name2 = generate(preprocessed) # Увеличение разрешения входного изображения
-    return image_pil, mesh_name, mesh_name2
-
-with gr.Blocks() as demo:
-    gr.Markdown("# **Input Image**")
-    with gr.Column():
-        input_image = gr.Image(
-            label="",
-            image_mode="RGBA",
-            sources="upload",
-            type="pil",
-            elem_id="content_image",
-        ) # Увеличение разрешения входного изображения
+    mesh_name, mesh_name2 = generate(preprocessed)
+    return preprocessed, mesh_name, mesh_name2
+
+with gr.Blocks(theme=CustomTheme(), css=css) as demo:
+    # **Header**
+    with gr.Row():
+        gr.Markdown("# 3D Model Generator", elem_id="title")
+
+    # **Input Section**
+    with gr.Accordion("Input Image", open=False):
         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("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.Tab("obj"):
-            output_model = gr.Model3D(
-                label="Output Model", 
-                interactive=False,
+            input_image = gr.Image(
+                label="Upload Image",
+                image_mode="RGBA",
+                sources="upload",
+                type="pil",
+                elem_id="content_image",
+                width=500,
             )
-        with gr.Tab("glb"):
-            output_model2 = gr.Model3D(
-                label="Output Model", 
-                interactive=False,
-            )
-    #...
-
-    with gr.Column():
-        gr.Examples(
-            examples=["examples/1.png"],
-            inputs=[input_image],
-            outputs=[input_image, output_model, output_model2],
-            fn=run_example,
-            cache_examples=True,
-        )
+            with gr.Row():
+                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,
+                )
+
+    # **Processing and Generation**
+    with gr.Accordion("Generate 3D Model", open=False):
+        with gr.Column():
+            submit = gr.Button("Generate", elem_classes="generate-button")
+            processed_image = gr.Image(label="Processed Image", interactive=False)
+            with gr.Tabs():
+                with gr.Tab("OBJ Model"):
+                    output_model = gr.Model3D(label="Output Model", interactive=False)
+                with gr.Tab("GLB Model"):
+                    output_model2 = gr.Model3D(label="Output Model", interactive=False)
+
+    # **Event Triggers**
+    submit.click(
+        fn=check_input_image, 
+        inputs=[input_image], 
+        outputs=[gr.update(label="Processing...", interactive=False)]
+    ).then(
+        fn=preprocess, 
+        inputs=[input_image, do_remove_background, foreground_ratio], 
+        outputs=[processed_image]
+    ).then(
+        fn=generate, 
+        inputs=[processed_image], 
+        outputs=[output_model, output_model2]
+    )
 
 demo.queue(max_size=10)
 demo.launch()