Create app.py

app.py

@@ -0,0 +1,257 @@
+import gradio as gr
+
+from canvas_editor import (
+    HTML_TEMPLATE,
+    CSS_TEMPLATE,
+    JS_ON_LOAD,
+    payload_append,
+    payload_reset,
+    empty_canvas_payload,
+)
+from shape_e_service import (
+    DEVICE,
+    EXAMPLES,
+    clear_saved_assets,
+    generate_and_add_asset,
+    next_view,
+    prev_view,
+    select_asset,
+    selected_view_path,
+    set_prompt,
+)
+
+CSS = """
+#title-wrap {
+    text-align: center;
+    margin-bottom: 10px;
+}
+#subtitle {
+    color: #9ca3af;
+    font-size: 14px;
+}
+#device-badge {
+    display: inline-block;
+    margin-top: 8px;
+    padding: 6px 10px;
+    border-radius: 999px;
+    background: #111827;
+    border: 1px solid #374151;
+    font-size: 13px;
+}
+"""
+
+
+def add_selected_view_to_canvas(saved_assets, selected_asset_index):
+    path = selected_view_path(saved_assets, selected_asset_index)
+    if not path:
+        raise gr.Error("Select an asset first.")
+    return payload_append([path])
+
+
+def reset_canvas_with_selected_view(saved_assets, selected_asset_index):
+    path = selected_view_path(saved_assets, selected_asset_index)
+    if not path:
+        raise gr.Error("Select an asset first.")
+    return payload_reset([path])
+
+
+def clear_canvas():
+    return empty_canvas_payload()
+
+
+with gr.Blocks(css=CSS, title="Depthpaint Studio") as demo:
+    saved_assets_state = gr.State([])
+    selected_asset_index_state = gr.State(None)
+
+    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="device-badge">Device: {DEVICE}</div>
+        </div>
+        """
+    )
+
+    with gr.Tab("Assets"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                prompt = gr.Textbox(
+                    label="Prompt",
+                    lines=3,
+                    value="A fantasy treasure chest with gold trim",
+                    placeholder="Describe a 3D object...",
+                )
+
+                gr.Markdown("**Prompt ideas**")
+                with gr.Row():
+                    ex1 = gr.Button("Robot")
+                    ex2 = gr.Button("Treasure chest")
+                    ex3 = gr.Button("Dragon toy")
+                with gr.Row():
+                    ex4 = gr.Button("Medieval house")
+                    ex5 = gr.Button("Owl teapot")
+                    ex6 = gr.Button("Submarine")
+
+                with gr.Row():
+                    steps = gr.Slider(
+                        minimum=8,
+                        maximum=64,
+                        step=1,
+                        value=24 if DEVICE == "cpu" else 32,
+                        label="Inference steps",
+                    )
+                    guidance_scale = gr.Slider(
+                        minimum=1.0,
+                        maximum=20.0,
+                        step=0.5,
+                        value=12.0,
+                        label="Guidance scale",
+                    )
+
+                with gr.Row():
+                    frame_size = gr.Slider(
+                        minimum=64,
+                        maximum=256,
+                        step=32,
+                        value=64 if DEVICE == "cpu" else 256,
+                        label="Frame size",
+                    )
+                    seed = gr.Number(
+                        label="Seed",
+                        value=42,
+                        precision=0,
+                    )
+
+                with gr.Row():
+                    generate_btn = gr.Button("Generate and add to gallery", variant="primary")
+                    clear_saved_btn = gr.Button("Clear gallery")
+
+                gr.Markdown(
+                    "CPU mode is much slower. Start with smaller frame size and fewer steps."
+                )
+
+            with gr.Column(scale=1):
+                current_view = gr.Image(label="Selected asset view", type="filepath")
+                view_text = gr.Markdown("No asset selected.")
+
+                with gr.Row():
+                    prev_btn = gr.Button("← Prev view", interactive=False)
+                    next_btn = gr.Button("Next view →", interactive=False)
+
+                with gr.Row():
+                    add_to_canvas_btn = gr.Button("Add selected view to canvas", variant="primary")
+                    reset_canvas_btn = gr.Button("Reset canvas with selected view")
+
+        saved_gallery = gr.Gallery(
+            label="Saved assets",
+            columns=3,
+            height="auto",
+            preview=False,
+        )
+
+    with gr.Tab("Canvas"):
+        gr.Markdown("## Scene canvas")
+        gr.Markdown("Selected asset views can be sent here and placed manually.")
+
+        editor = gr.HTML(
+            value='{"render_id": null, "mode": "append", "items": []}',
+            html_template=HTML_TEMPLATE,
+            css_template=CSS_TEMPLATE,
+            js_on_load=JS_ON_LOAD,
+        )
+
+        clear_canvas_btn = gr.Button("Clear canvas")
+
+    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)
+    ex4.click(fn=lambda: set_prompt(EXAMPLES[3]), outputs=prompt)
+    ex5.click(fn=lambda: set_prompt(EXAMPLES[4]), outputs=prompt)
+    ex6.click(fn=lambda: set_prompt(EXAMPLES[5]), outputs=prompt)
+
+    generate_btn.click(
+        fn=generate_and_add_asset,
+        inputs=[prompt, steps, guidance_scale, frame_size, seed, saved_assets_state],
+        outputs=[
+            saved_assets_state,
+            selected_asset_index_state,
+            saved_gallery,
+            current_view,
+            view_text,
+            prev_btn,
+            next_btn,
+        ],
+    )
+
+    saved_gallery.select(
+        fn=select_asset,
+        inputs=[saved_assets_state],
+        outputs=[
+            selected_asset_index_state,
+            saved_gallery,
+            current_view,
+            view_text,
+            prev_btn,
+            next_btn,
+        ],
+    )
+
+    prev_btn.click(
+        fn=prev_view,
+        inputs=[saved_assets_state, selected_asset_index_state],
+        outputs=[
+            saved_assets_state,
+            saved_gallery,
+            current_view,
+            view_text,
+            prev_btn,
+            next_btn,
+        ],
+    )
+
+    next_btn.click(
+        fn=next_view,
+        inputs=[saved_assets_state, selected_asset_index_state],
+        outputs=[
+            saved_assets_state,
+            saved_gallery,
+            current_view,
+            view_text,
+            prev_btn,
+            next_btn,
+        ],
+    )
+
+    clear_saved_btn.click(
+        fn=clear_saved_assets,
+        outputs=[
+            saved_assets_state,
+            selected_asset_index_state,
+            saved_gallery,
+            current_view,
+            view_text,
+            prev_btn,
+            next_btn,
+        ],
+    )
+
+    add_to_canvas_btn.click(
+        fn=add_selected_view_to_canvas,
+        inputs=[saved_assets_state, selected_asset_index_state],
+        outputs=editor,
+    )
+
+    reset_canvas_btn.click(
+        fn=reset_canvas_with_selected_view,
+        inputs=[saved_assets_state, selected_asset_index_state],
+        outputs=editor,
+    )
+
+    clear_canvas_btn.click(
+        fn=clear_canvas,
+        outputs=editor,
+    )
+
+if __name__ == "__main__":
+    demo.launch()