Update app.py

app.py

@@ -1,9 +1,4 @@
 from pathlib import Path
-import gradio as gr
-
-ROOT_DIR = Path(__file__).resolve().parent
-DATA_DIR = ROOT_DIR / "data"
-DATA_DIR.mkdir(parents=True, exist_ok=True)
 
 import gradio as gr
 
@@ -26,6 +21,11 @@ from shape_e_service import (
     selected_view_path,
     set_prompt,
 )
+from sd_service import generate_with_depth_from_scene
+
+ROOT_DIR = Path(__file__).resolve().parent
+DATA_DIR = ROOT_DIR / "data"
+DATA_DIR.mkdir(parents=True, exist_ok=True)
 
 CSS = """
 #title-wrap {
@@ -45,6 +45,9 @@ CSS = """
     border: 1px solid #374151;
     font-size: 13px;
 }
+#scene_png_data {
+    display: none;
+}
 """
 
 
@@ -66,15 +69,18 @@ def clear_canvas():
     return empty_canvas_payload()
 
 
-with gr.Blocks(css=CSS, title="Depthpaint Studio") as demo:
+with gr.Blocks(title="Depthpaint Studio") as demo:
     saved_assets_state = gr.State([])
     selected_asset_index_state = gr.State(None)
 
+    # Скрытое поле, в которое canvas JS складывает PNG композиции в виде data URL
+    scene_png_data = gr.Textbox(elem_id="scene_png_data", value="", lines=1)
+
     gr.HTML(
         f"""
         <div id="title-wrap">
             <h1>Depthpaint Studio</h1>
-            <div id="subtitle">Generate Shap-E assets, browse their views, then send the selected view into the canvas</div>
+            <div id="subtitle">Generate Shap-E assets, browse their views, place them on the canvas, then render with SD + ControlNet depth</div>
             <div id="device-badge">Device: {DEVICE}</div>
         </div>
         """
@@ -168,6 +174,40 @@ with gr.Blocks(css=CSS, title="Depthpaint Studio") as demo:
 
     clear_canvas_btn = gr.Button("Clear canvas")
 
+    gr.Markdown("## SD + ControlNet depth")
+    gr.Markdown("Current canvas composition will be sent into SD as init image, and its depth will be used as ControlNet guidance.")
+
+    sd_prompt = gr.Textbox(
+        label="SD prompt",
+        lines=3,
+        value="A cozy fantasy tabletop scene, detailed, cinematic lighting, highly detailed",
+    )
+
+    sd_negative_prompt = gr.Textbox(
+        label="Negative prompt",
+        lines=2,
+        value="low quality, blurry, cropped, deformed, duplicate objects, floating objects, white background",
+    )
+
+    with gr.Row():
+        sd_steps = gr.Slider(label="SD steps", minimum=10, maximum=50, step=1, value=30)
+        sd_guidance = gr.Slider(label="Guidance scale", minimum=1.0, maximum=12.0, step=0.5, value=7.5)
+
+    with gr.Row():
+        sd_strength = gr.Slider(label="Img2Img strength", minimum=0.1, maximum=1.0, step=0.05, value=0.65)
+        controlnet_scale = gr.Slider(label="Depth control scale", minimum=0.1, maximum=2.0, step=0.1, value=0.9)
+
+    sd_seed = gr.Number(label="SD seed", value=1234, precision=0)
+
+    render_sd_btn = gr.Button("Render canvas with SD + depth", variant="primary")
+
+    with gr.Row():
+        scene_preview = gr.Image(label="Canvas composition sent to SD", type="filepath")
+        depth_preview = gr.Image(label="Depth used by ControlNet", type="filepath")
+
+    sd_result = gr.Image(label="SD result", type="filepath")
+
+    # Prompt example buttons
     ex1.click(fn=lambda: set_prompt(EXAMPLES[0]), outputs=prompt)
     ex2.click(fn=lambda: set_prompt(EXAMPLES[1]), outputs=prompt)
     ex3.click(fn=lambda: set_prompt(EXAMPLES[2]), outputs=prompt)
@@ -175,6 +215,7 @@ with gr.Blocks(css=CSS, title="Depthpaint Studio") as demo:
     ex5.click(fn=lambda: set_prompt(EXAMPLES[4]), outputs=prompt)
     ex6.click(fn=lambda: set_prompt(EXAMPLES[5]), outputs=prompt)
 
+    # Shape-E generation
     generate_btn.click(
         fn=generate_and_add_asset,
         inputs=[prompt, steps, guidance_scale, frame_size, seed, saved_assets_state],
@@ -189,6 +230,7 @@ with gr.Blocks(css=CSS, title="Depthpaint Studio") as demo:
         ],
     )
 
+    # Gallery selection
     saved_gallery.select(
         fn=select_asset,
         inputs=[saved_assets_state],
@@ -202,6 +244,7 @@ with gr.Blocks(css=CSS, title="Depthpaint Studio") as demo:
         ],
     )
 
+    # View switching
     prev_btn.click(
         fn=prev_view,
         inputs=[saved_assets_state, selected_asset_index_state],
@@ -228,6 +271,7 @@ with gr.Blocks(css=CSS, title="Depthpaint Studio") as demo:
         ],
     )
 
+    # Clear gallery
     clear_saved_btn.click(
         fn=clear_saved_assets,
         outputs=[
@@ -241,6 +285,7 @@ with gr.Blocks(css=CSS, title="Depthpaint Studio") as demo:
         ],
     )
 
+    # Canvas actions
     add_to_canvas_btn.click(
         fn=add_selected_view_to_canvas,
         inputs=[saved_assets_state, selected_asset_index_state],
@@ -258,5 +303,24 @@ with gr.Blocks(css=CSS, title="Depthpaint Studio") as demo:
         outputs=editor,
     )
 
+    # SD render
+    render_sd_btn.click(
+        fn=generate_with_depth_from_scene,
+        inputs=[
+            scene_png_data,
+            sd_prompt,
+            sd_negative_prompt,
+            sd_steps,
+            sd_guidance,
+            sd_strength,
+            controlnet_scale,
+            sd_seed,
+        ],
+        outputs=[scene_preview, depth_preview, sd_result],
+    )
+
 if __name__ == "__main__":
-    demo.launch(allowed_paths=[str(DATA_DIR)])
+    demo.launch(
+        css=CSS,
+        allowed_paths=[str(DATA_DIR)],
+    )