Update app.py

app.py

@@ -1,27 +1,16 @@
-import logging
-import os
-import tempfile
-import time
-
 import gradio as gr
+import tempfile
+import os
 import numpy as np
-import rembg
 import torch
 from PIL import Image
-from functools import partial
-
 from tsr.system import TSR
 from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation
 
-if torch.cuda.is_available():
-    device = "cuda:0"
-else:
-    device = "cpu"
-
-d = os.environ.get("DEVICE", None)
-if d != None:
-    device = d 
+# Проверяем наличие GPU
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
+# Загружаем модель
 model = TSR.from_pretrained(
     "stabilityai/TripoSR",
     config_name="config.yaml",
@@ -30,14 +19,14 @@ model = TSR.from_pretrained(
 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
@@ -47,7 +36,7 @@ def preprocess(input_image, do_remove_background, foreground_ratio):
 
     if do_remove_background:
         image = input_image.convert("RGB")
-        image = remove_background(image, rembg_session)
+        image = remove_background(image)
         image = resize_foreground(image, foreground_ratio)
         image = fill_background(image)
     else:
@@ -57,6 +46,7 @@ def preprocess(input_image, do_remove_background, foreground_ratio):
     return image
 
 
+# Функция генерации 3D модели
 def generate(image):
     scene_codes = model(image, device=device)
     mesh = model.extract_mesh(scene_codes)[0]
@@ -66,6 +56,7 @@ def generate(image):
     return mesh_path2.name
 
 
+# Настройка темы и CSS
 class CustomTheme(gr.themes.Base):
     def __init__(self):
         super().__init__()
@@ -81,145 +72,64 @@ class CustomTheme(gr.themes.Base):
 
 
 css = """
-footer {
-    visibility: hidden;
-    height: 0;
-    margin: 0;
-    padding: 0;
-    overflow: hidden;
-}
-
-body, input, button, textarea, select, .gr-button {
-    font-family: 'Poppins', sans-serif;
+body {
     background-color: #191a1e !important;
-    color: #FFFFFF;
 }
 
-h1, h2, h3, h4, h5, h6 {
-    font-family: 'Poppins', sans-serif;
-    font-weight: 700;
+.gr-container {
+    background-color: #191a1e !important;
     color: #FFFFFF;
 }
 
-input[type="text"], textarea {
+.gr-box, .gr-button, input, textarea {
     background-color: #191a1e !important;
-    color: #FFFFFF;
+    color: #FFFFFF !important;
     border: 1px solid #FFFFFF;
 }
 
-.generate-button {
+.gr-button {
     background-color: #5271FF !important;
-    color: #FFFFFF !important;
-    border: none;
-    font-weight: bold;
 }
 
-.generate-button:hover {
+.gr-button:hover {
     background-color: #405BBF !important;
 }
-
-gr-image-upload {
-    background-color: #191a1e !important;
-    border: 1px dashed #FFFFFF !important;
-    position: relative;
-}
-
-.gr-image-upload .dropzone::before {
-    content: 'Drop Image Here';
-    color: #FFFFFF;
-    font-size: 1rem;
-    font-weight: bold;
-    text-align: center;
-    position: absolute;
-    top: 50%;
-    left: 50%;
-    transform: translate(-50%, -50%);
-}
-
-.gr-image-upload .dropzone .file-upload {
-    display: none !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;
-}
-
-/* Вкладки GLB Model */
-.gr-tab {
-    background-color: #252525; /* Чуть светлее общий фон */
-    border: 1px solid #FFFFFF; /* Белая рамка для структуры */
-    border-radius: 5px; /* Скругление рамок */
-    margin-top: 10px; /* Отступ от верхнего элемента */
-}
-
-.gr-tab label {
-    color: #FFFFFF; /* Цвет текста */
-    font-weight: bold;
-}
-
 """
 
+# Интерфейс
 with gr.Blocks(theme=CustomTheme(), css=css) as demo:
-    with gr.Column(elem_id="main-container"):
-        gr.Markdown("**Upload and Process Your Image**", elem_classes="prompt-text")
-
-        with gr.Row(elem_id="image-row"):
+    with gr.Column():
+        gr.Markdown("**Upload and Process Your Image**")
+        with gr.Row():
             input_image = gr.Image(
                 label="Upload Image",
                 image_mode="RGBA",
                 sources="upload",
                 type="pil",
-                elem_id="content_image",
                 width=400,
                 height=300,
             )
             processed_image = gr.Image(
                 label="Processed Image",
                 interactive=False,
-                elem_id="processed_image",
                 width=400,
                 height=300,
             )
+        foreground_ratio = gr.Slider(
+            label="Foreground Ratio",
+            minimum=0.5,
+            maximum=1.0,
+            value=0.85,
+            step=0.05,
+        )
+        do_remove_background = gr.Checkbox(
+            label="Remove Background",
+            value=True,
+        )
+        submit = gr.Button("Generate")
+        output_model = gr.Model3D(label="Generated GLB Model", interactive=False)
 
-        with gr.Row(elem_id="slider-row"):
-            foreground_ratio = gr.Slider(
-                label="Foreground Ratio",
-                minimum=0.5,
-                maximum=1.0,
-                value=0.85,
-                step=0.05,
-            )
-
-        with gr.Row(elem_id="settings-row"):
-            do_remove_background = gr.Checkbox(
-                label="Remove Background",
-                value=True,
-            )
-            submit = gr.Button(
-                "Generate",
-                scale=0,
-                variant="primary",
-                elem_classes="generate-button",
-            )
-
-        with gr.Row(elem_id="output-model-container"):
-            with gr.Tab("GLB Model"):
-                output_model2 = gr.Model3D(
-                    label="Generated GLB Model",
-                    interactive=False,
-                )
-
+        # Обработка кликов
         submit.click(
             fn=check_input_image,
             inputs=[input_image],
@@ -231,9 +141,8 @@ with gr.Blocks(theme=CustomTheme(), css=css) as demo:
         ).then(
             fn=generate,
             inputs=[processed_image],
-            outputs=[output_model2]
+            outputs=[output_model]
         )
 
-
-demo.queue(max_size=10)
+# Запуск приложения
 demo.launch()