Upload app.py

app.py

@@ -15,6 +15,12 @@ from transformers import CLIPTokenizer, CLIPTextModel, CLIPImageProcessor
 import torch.nn as nn
 import numpy as np
 
+import tempfile
+import os
+import shutil
+import uuid
+
+
 class UNetNoCondWrapper(nn.Module):
     def __init__(self, base_unet: UNet2DModel):
         super().__init__()
@@ -71,243 +77,260 @@ class UNetNoCondWrapper(nn.Module):
 # pipe = pipe.to(torch.float32).to(device)
 # @spaces.GPU
 
-def pil_to_data_uri(img: Image.Image) -> str:
-    buf = io.BytesIO()
-    img.save(buf, format="PNG")
-    b = base64.b64encode(buf.getvalue()).decode("utf-8")
-    return f"data:image/png;base64,{b}"
-
-# # --- new gradio_generate that returns both gallery and 3D html ---
-# def gradio_generate(fibers: Image.Image, rings: Image.Image, num_steps: int):
-#     """
-#     fibers: PIL image containing 4 tiles in a 2x2 square (order 1,4,2,3)
-#     rings:  same
-#     returns: (list_of_4_PIL_images, html_string_for_3d)
-#     """
-#     # call your inference function (unchanged)
-#     preds = inference(pipe, fibers, rings, int(num_steps))  # preds order: [1,4,2,3] -> TL,TR,BL,BR
-
-#     # convert preds to RGB for gallery
-#     preds_rgb = [p.convert("RGB") if p.mode != "RGB" else p for p in preds]
-
-#     # build data-urls in same order
-#     data_urls = [pil_to_dataurl(im, fmt="PNG") for im in preds_rgb]  # [TL,TR,BL,BR]
-
-#     # safe JSON array literal for injection into HTML template
-#     js_img_list = json.dumps(data_urls)
-
-#     html = HTML_TEMPLATE.replace("___IMGS___", js_img_list)
-
-#     # gr.Interface expects outputs matching signature; we return (gallery_images, html_string)
-#     return preds_rgb, html
-
-# # replace the Interface outputs: first a Gallery, second a HTML component
-# iface = gr.Interface(
-#     fn=gradio_generate,
-#     inputs=[
-#         gr.Image(type="pil", label="Fiber"),
-#         gr.Image(type="pil", label="Ring"),
-#         gr.Number(value=10, label="Number of inference steps"),
-#     ],
-#     outputs=[
-#         gr.HTML(label="3D preview")
-#     ],
-#     title="Photorealistic wood generator (4 faces)",
-#     description="""
-# Upload 2 images (one with four fiber maps and one with four ring maps) arranged as a 2x2 square.
-# The model returns four generated faces (one per tile) and a 3D preview with those faces textured on a box.
-# """
-# )
+# def pil_to_data_uri(img: Image.Image) -> str:
+#     buf = io.BytesIO()
+#     img.save(buf, format="PNG")
+#     b = base64.b64encode(buf.getvalue()).decode("utf-8")
+#     return f"data:image/png;base64,{b}"
 
-def run(fibers: Image.Image, rings: Image.Image, num_steps: int):
+def user_inference(img1, img2):
+    try:
+        # si numpy array -> PIL
+        if isinstance(img1, np.ndarray):
+            img1 = Image.fromarray(img1)
+        if isinstance(img2, np.ndarray):
+            img2 = Image.fromarray(img2)
+        # 4 outputs tests
+        o1 = img1.copy().resize((512,512))
+        o2 = img2.copy().resize((512,512))
+        o3 = img1.transpose(Image.FLIP_LEFT_RIGHT).resize((512,512))
+        o4 = img2.rotate(90, expand=True).resize((512,512))
+        return [o1, o2, o3, o4]
+    except Exception:
+        blank = Image.new("RGB", (512,512), (200,200,200))
+        return [blank, blank, blank, blank]
+    
+def build_textured_cube(textures, out_dir=None):
+    """
+    textures: list de 4 PIL.Image (RGB)
+    out_dir: optionnel dossier de sortie ; si None, un tmpdir est créé
+    retourne: chemin absolu vers le fichier .obj
+    """
+    if out_dir is None:
+        out_dir = tempfile.mkdtemp(prefix="cube_")
+    else:
+        os.makedirs(out_dir, exist_ok=True)
+
+    # sauvegarde textures
+    tex_filenames = []
+    for i, im in enumerate(textures):
+        fn = f"tex_{i}.png"
+        path = os.path.join(out_dir, fn)
+        im.save(path, format="PNG")
+        tex_filenames.append(fn)
+
+    # mtl content (référence les textures sauvegardées)
+    mtl_name = "cube.mtl"
+    mtl_path = os.path.join(out_dir, mtl_name)
+    with open(mtl_path, "w") as f:
+        # materials mapping: mat_right, mat_left, mat_top, mat_bottom, mat_front, mat_back
+        f.write("# materials for cube\n")
+        # right -> textures[0]
+        f.write("newmtl mat_right\n")
+        f.write(f"  Ns 10\n")
+        f.write(f"  Ka 1.000 1.000 1.000\n")
+        f.write(f"  Kd 1.000 1.000 1.000\n")
+        f.write(f"  map_Kd {tex_filenames[0]}\n\n")
+        # left -> textures[1]
+        f.write("newmtl mat_left\n")
+        f.write("  Ns 10\n")
+        f.write("  Ka 1.000 1.000 1.000\n")
+        f.write("  Kd 1.000 1.000 1.000\n")
+        f.write(f"  map_Kd {tex_filenames[1]}\n\n")
+        # top neutral
+        f.write("newmtl mat_top\n")
+        f.write("  Ns 10\n")
+        f.write("  Ka 1.000 1.000 1.000\n")
+        f.write("  Kd 0.8 0.8 0.8\n\n")
+        # bottom neutral
+        f.write("newmtl mat_bottom\n")
+        f.write("  Ns 10\n")
+        f.write("  Ka 1.000 1.000 1.000\n")
+        f.write("  Kd 0.7 0.7 0.7\n\n")
+        # front -> textures[2]
+        f.write("newmtl mat_front\n")
+        f.write("  Ns 10\n")
+        f.write("  Ka 1.000 1.000 1.000\n")
+        f.write("  Kd 1.000 1.000 1.000\n")
+        f.write(f"  map_Kd {tex_filenames[2]}\n\n")
+        # back -> textures[3]
+        f.write("newmtl mat_back\n")
+        f.write("  Ns 10\n")
+        f.write("  Ka 1.000 1.000 1.000\n")
+        f.write("  Kd 1.000 1.000 1.000\n")
+        f.write(f"  map_Kd {tex_filenames[3]}\n\n")
+
+    # OBJ content
+    obj_name = "cube.obj"
+    obj_path = os.path.join(out_dir, obj_name)
+    with open(obj_path, "w") as f:
+        f.write(f"mtllib {mtl_name}\n")
+        f.write("o CubeTextured\n")
+
+        # 8 vertices
+        verts = [
+            (-1.0, -1.0, -1.0),  # v1
+            ( 1.0, -1.0, -1.0),  # v2
+            ( 1.0,  1.0, -1.0),  # v3
+            (-1.0,  1.0, -1.0),  # v4
+            (-1.0, -1.0,  1.0),  # v5
+            ( 1.0, -1.0,  1.0),  # v6
+            ( 1.0,  1.0,  1.0),  # v7
+            (-1.0,  1.0,  1.0),  # v8
+        ]
+        for v in verts:
+            f.write(f"v {v[0]} {v[1]} {v[2]}\n")
+
+        # Pour chaque face on va définir 4 UVs (vt) — on crée 4 vt par face (24 au total)
+        # uv order: (0,0),(1,0),(1,1),(0,1)
+        uv_coords = [(0.0,0.0),(1.0,0.0),(1.0,1.0),(0.0,1.0)]
+        for _ in range(6):
+            for uv in uv_coords:
+                f.write(f"vt {uv[0]} {uv[1]}\n")
+
+        # Faces:
+        # We'll use materials in this order: right, left, top, bottom, front, back
+        # Note: vt indices are 1..24
+        # indices for convenience
+        # vertices indices: v1..v8 (1-based)
+        # define faces as quads (OBJ supports quads)
+        # right  (x = +1): v2, v6, v7, v3
+        f.write("usemtl mat_right\n")
+        f.write("f 2/1 6/2 7/3 3/4\n")
+        # left (x = -1): v5, v1, v4, v8
+        f.write("usemtl mat_left\n")
+        f.write("f 5/5 1/6 4/7 8/8\n")
+        # top (y = +1): v4, v3, v7, v8
+        f.write("usemtl mat_top\n")
+        f.write("f 4/9 3/10 7/11 8/12\n")
+        # bottom (y = -1): v1, v5, v6, v2
+        f.write("usemtl mat_bottom\n")
+        f.write("f 1/13 5/14 6/15 2/16\n")
+        # front (z = +1): v5, v6, v7, v8
+        f.write("usemtl mat_front\n")
+        f.write("f 5/17 6/18 7/19 8/20\n")
+        # back (z = -1): v1, v2, v3, v4
+        f.write("usemtl mat_back\n")
+        f.write("f 1/21 2/22 3/23 4/24\n")
+
+    # return absolute path to .obj
+    return obj_path, out_dir
+
+# -------------------------
+# Fonction appelée par Gradio
+# retourne : 4 miniatures (PIL) + chemin vers le .obj (str)
+# -------------------------
+def run(img1, img2):
     try:
-        # appelle la fonction d'inference de l'utilisateur
-        # outputs = inference(pipe, fibers, rings, int(num_steps))
-        outputs = inference(fibers, rings, int(num_steps))
-        if not (isinstance(outputs, (list,tuple)) and len(outputs) >= 4):
-            raise ValueError("La fonction d'inference doit renvoyer une liste/tuple de 4 images.")
-        # Prendre les 4 premières images
-        imgs = outputs[:4]
-        # Convertir en PIL si nécessaire et assurer taille raisonnable
+        outputs = user_inference(img1, img2)
+        if not (isinstance(outputs, (list, tuple)) and len(outputs) >= 4):
+            raise ValueError("user_inference doit renvoyer une liste/tuple de 4 images.")
+        # Prendre les 4 premières images et convertir en PIL RGB
         pil_imgs = []
-        for im in imgs:
+        for im in outputs[:4]:
             if isinstance(im, np.ndarray):
                 im = Image.fromarray(im)
-            # for safety, ensure RGB
             if im.mode != "RGB":
                 im = im.convert("RGB")
             pil_imgs.append(im)
 
-        # Préparer miniatures (pour les 4 gr.Image)
-        thumb_imgs = [im.copy().resize((256,256)) for im in pil_imgs]
-
-        # Convertir en data-URIs pour Three.js
-        data_uris = [pil_to_data_uri(im) for im in pil_imgs]
-
-        # Construire le HTML/JS pour le visualiseur Three.js
-        html = f"""
-<div id="viewer" style="width:100%;height:480px; border:1px solid #ddd; position:relative;background:#f6f6f6;">
-  <div id="viewer-msg" style="position:absolute;left:8px;top:8px;padding:6px 8px;background:rgba(255,255,255,0.9);border-radius:6px;font-size:12px;color:#333;z-index:2;">
-    Manipulez la souris pour tourner, molette pour zoomer.
-  </div>
-</div>
-<script src="https://unpkg.com/three@0.152.2/build/three.min.js"></script>
-<script src="https://unpkg.com/three@0.152.2/examples/js/controls/OrbitControls.js"></script>
-<script>
-(function() {{
-  const images = {data_uris!r};
-  const container = document.getElementById('viewer');
-
-  function safeLog(...args) {{ try{{console.log(...args)}}catch(e){{}} }}
-
-  function initScene() {{
-    const w = container.clientWidth || 800;
-    const h = container.clientHeight || 480;
-
-    const scene = new THREE.Scene();
-    const camera = new THREE.PerspectiveCamera(45, w/h, 0.1, 1000);
-    camera.position.set(2.5, 2.0, 3.5);
-
-    const renderer = new THREE.WebGLRenderer({{antialias:true}});
-    renderer.setSize(w, h);
-    renderer.domElement.style.width = '100%';
-    renderer.domElement.style.height = '100%';
-    container.appendChild(renderer.domElement);
-
-    // lights
-    const ambient = new THREE.AmbientLight(0xffffff, 0.7);
-    scene.add(ambient);
-    const dir = new THREE.DirectionalLight(0xffffff, 0.6);
-    dir.position.set(5,10,7);
-    scene.add(dir);
-
-    // fallback cube (visible immédiatement)
-    const geometry = new THREE.BoxGeometry(1.6,1.6,1.6);
-    const fallbackMats = [];
-    for (let i=0;i<6;i++) fallbackMats.push(new THREE.MeshStandardMaterial({{color:0x999999}}));
-    const cube = new THREE.Mesh(geometry, fallbackMats);
-    scene.add(cube);
-
-    // grid for context
-    const grid = new THREE.GridHelper(6, 12, 0x888888, 0x444444);
-    grid.position.y = -1.2;
-    scene.add(grid);
-
-    // controls
-    const controls = new THREE.OrbitControls(camera, renderer.domElement);
-    controls.enableDamping = true;
-    controls.dampingFactor = 0.08;
-    controls.target.set(0,0,0);
-
-    // attempt to load textures, but keep fallback if it fails
-    const loader = new THREE.TextureLoader();
-    Promise.all(images.map(uri => new Promise((res, rej) => {{
-      try {{
-        loader.load(uri, tex => {{ tex.flipY = false; tex.generateMipmaps = true; res(tex); }}, undefined, err => rej(err));
-      }} catch(e) {{ rej(e); }}
-    }}))).then(textures => {{
-      safeLog('Textures loaded', textures);
-      const mats = [
-        new THREE.MeshStandardMaterial({{map: textures[0]}}), // right
-        new THREE.MeshStandardMaterial({{map: textures[1]}}), // left
-        new THREE.MeshStandardMaterial({{color:0xcccccc}}),    // top
-        new THREE.MeshStandardMaterial({{color:0xcccccc}}),    // bottom
-        new THREE.MeshStandardMaterial({{map: textures[2]}}), // front
-        new THREE.MeshStandardMaterial({{map: textures[3]}})  // back
-      ];
-      // set wrapping/filter
-      mats.forEach(m => {{
-        if (m.map) {{
-          m.map.minFilter = THREE.LinearFilter;
-          m.map.wrapS = THREE.ClampToEdgeWrapping;
-          m.map.wrapT = THREE.ClampToEdgeWrapping;
-        }}
-      }});
-      cube.material = mats;
-    }}).catch(err => {{
-      safeLog('Erreur chargement textures (on garde fallback):', err);
-      // affiche un petit message utilisateur (optionnel)
-      const msg = document.getElementById('viewer-msg');
-      if (msg) msg.innerText = "Rendu 3D : textures non chargées, affichage fallback (voir console).";
-    }});
-
-    // animation loop
-    function animate() {{
-      requestAnimationFrame(animate);
-      controls.update();
-      renderer.render(scene, camera);
-    }}
-    animate();
-
-    // responsive: observe container size
-    if (typeof ResizeObserver !== 'undefined') {{
-      const ro = new ResizeObserver(() => {{
-        const ww = container.clientWidth || 800;
-        const hh = container.clientHeight || 480;
-        renderer.setSize(ww, hh);
-        camera.aspect = ww / hh;
-        camera.updateProjectionMatrix();
-      }});
-      ro.observe(container);
-    }} else {{
-      window.addEventListener('resize', () => {{
-        const ww = container.clientWidth || 800;
-        const hh = container.clientHeight || 480;
-        renderer.setSize(ww, hh);
-        camera.aspect = ww / hh;
-        camera.updateProjectionMatrix();
-      }});
-    }}
-  }}
-
-  // wait until container has non-zero size and THREE is available
-  let tries = 0;
-  function waitForReady() {{
-    tries++;
-    if (container.clientWidth > 0 && typeof THREE !== 'undefined') {{
-      try {{ initScene(); }} catch(e) {{ safeLog('Erreur initScene', e); container.innerHTML = "<div style='padding:10px;color:#900;'>Erreur initialisation rendu 3D (voir console)</div>"; }}
-    }} else {{
-      if (tries < 60) setTimeout(waitForReady, 100);
-      else {{
-        // give it one last try
-        try {{ initScene(); }} catch(e) {{ safeLog('Derniere tentative échouée', e); container.innerHTML = "<div style='padding:10px;color:#900;'>Impossible d'initialiser le rendu 3D (voir console)</div>"; }}
-      }}
-    }}
-  }}
-  waitForReady();
-}})();
-</script>
-"""
-        # Retourner : 4 miniatures pour affichage + HTML viewer
-        return (*thumb_imgs, html)
+        # miniatures affichées dans l'UI
+        thumbs = [im.copy().resize((256,256)) for im in pil_imgs]
+
+        # construire le .obj + textures
+        obj_path, tmpdir = build_textured_cube(pil_imgs)
+
+        # IMPORTANT: on renvoie le chemin absolu vers le fichier .obj
+        # Gradio/Spaces chargera le modèle 3D depuis ce fichier et les textures
+        # doivent être dans le même dossier que le .obj (ce qui est le cas).
+        return (*thumbs, obj_path)
     except Exception as e:
         traceback.print_exc()
-        # en cas d'erreur, renvoyer 4 images vides et message html d'erreur
+        # en cas d'erreur, renvoyer 4 images grises et None pour le modèle
         blank = Image.new("RGB", (256,256), (220,220,220))
-        err_html = f"<div style='color:#900;padding:12px;'>Erreur serveur: {str(e)}</div>"
-        return (blank, blank, blank, blank, err_html)
+        return (blank, blank, blank, blank, None)
 
+# -------------------------
 # Interface Gradio
-with gr.Blocks(title="3D Cube Viewer - 4 outputs on cube") as demo:
-    gr.Markdown("### App 3D : assemble 4 images produites par votre code d'inference sur les 4 faces d'un cube.")
+# -------------------------
+with gr.Blocks(title="Cube 3D — coller 4 images sur les faces") as demo:
+    gr.Markdown("**Chargez 2 images** → votre code d'inference doit produire 4 images → les 4 images seront collées sur les faces latérales d'un cube (.obj).")
     with gr.Row():
         with gr.Column(scale=1):
-            inp1 = gr.Image(label="Image 1 (entrée)", type="numpy")
-            inp2 = gr.Image(label="Image 2 (entrée)", type="numpy")
-            inp3 = gr.Number(value=10, label="Number of inference steps")
+            inp1 = gr.Image(label="Image 1 (entrée)", type="pil")
+            inp2 = gr.Image(label="Image 2 (entrée)", type="pil")
             run_btn = gr.Button("Lancer l'inference")
-            gr.Markdown("Remplace `user_inference` dans `app.py` par ton code d'inférence existant.")
+            gr.Markdown("Remplace `user_inference` dans `app.py` par ta fonction d'inference réelle.")
         with gr.Column(scale=2):
-            viewer = gr.HTML("<div style='padding:12px;color:#666;'>Le rendu 3D s'affichera ici après inference.</div>", label="Visu 3D")
-    # gr.Row().style(mobile_collapse=False)
+            model3d_out = gr.Model3D(label="Cube 3D (obj + mtl + textures)")
     with gr.Row():
-        out1 = gr.Image(label="Output 1")
-        out2 = gr.Image(label="Output 2")
-        out3 = gr.Image(label="Output 3")
-        out4 = gr.Image(label="Output 4")
+        out1 = gr.Image(label="Output 1 (right)")
+        out2 = gr.Image(label="Output 2 (left)")
+        out3 = gr.Image(label="Output 3 (front)")
+        out4 = gr.Image(label="Output 4 (back)")
 
-    run_btn.click(fn=run, inputs=[inp1, inp2, inp3], outputs=[out1, out2, out3, out4, viewer])
+    run_btn.click(fn=run, inputs=[inp1, inp2], outputs=[out1, out2, out3, out4, model3d_out])
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
+    demo.launch(server_name="0.0.0.0", server_port=7860)
+
+# def run(fibers: Image.Image, rings: Image.Image, num_steps: int):
+#     try:
+#         # appelle la fonction d'inference de l'utilisateur
+#         # outputs = inference(pipe, fibers, rings, int(num_steps))
+#         outputs = inference(fibers, rings, int(num_steps))
+#         if not (isinstance(outputs, (list,tuple)) and len(outputs) >= 4):
+#             raise ValueError("La fonction d'inference doit renvoyer une liste/tuple de 4 images.")
+#         # Prendre les 4 premières images
+#         imgs = outputs[:4]
+#         # Convertir en PIL si nécessaire et assurer taille raisonnable
+#         pil_imgs = []
+#         for im in imgs:
+#             if isinstance(im, np.ndarray):
+#                 im = Image.fromarray(im)
+#             # for safety, ensure RGB
+#             if im.mode != "RGB":
+#                 im = im.convert("RGB")
+#             pil_imgs.append(im)
+
+#         # Préparer miniatures (pour les 4 gr.Image)
+#         thumb_imgs = [im.copy().resize((256,256)) for im in pil_imgs]
+
+#         # Convertir en data-URIs pour Three.js
+#         data_uris = [pil_to_data_uri(im) for im in pil_imgs]
+
+#         # Construire le HTML/JS pour le visualiseur Three.j
+#         html = 0
+#         # Retourner : 4 miniatures pour affichage + HTML viewer
+#         return (*thumb_imgs, html)
+#     except Exception as e:
+#         traceback.print_exc()
+#         # en cas d'erreur, renvoyer 4 images vides et message html d'erreur
+#         blank = Image.new("RGB", (256,256), (220,220,220))
+#         err_html = f"<div style='color:#900;padding:12px;'>Erreur serveur: {str(e)}</div>"
+#         return (blank, blank, blank, blank, err_html)
+
+# # Interface Gradio
+# with gr.Blocks(title="3D Cube Viewer - 4 outputs on cube") as demo:
+#     gr.Markdown("### App 3D : assemble 4 images produites par votre code d'inference sur les 4 faces d'un cube.")
+#     with gr.Row():
+#         with gr.Column(scale=1):
+#             inp1 = gr.Image(label="Image 1 (entrée)", type="numpy")
+#             inp2 = gr.Image(label="Image 2 (entrée)", type="numpy")
+#             inp3 = gr.Number(value=10, label="Number of inference steps")
+#             run_btn = gr.Button("Lancer l'inference")
+#             gr.Markdown("Remplace `user_inference` dans `app.py` par ton code d'inférence existant.")
+#         with gr.Column(scale=2):
+#             viewer = gr.HTML("<div style='padding:12px;color:#666;'>Le rendu 3D s'affichera ici après inference.</div>", label="Visu 3D")
+#     # gr.Row().style(mobile_collapse=False)
+#     with gr.Row():
+#         out1 = gr.Image(label="Output 1")
+#         out2 = gr.Image(label="Output 2")
+#         out3 = gr.Image(label="Output 3")
+#         out4 = gr.Image(label="Output 4")
+
+#     run_btn.click(fn=run, inputs=[inp1, inp2, inp3], outputs=[out1, out2, out3, out4, viewer])
+
+# if __name__ == "__main__":
+#     demo.launch(server_name="0.0.0.0", server_port=7860, share=True)