- """
-
- return image_tag
diff --git a/craftsman/utils/visualizers/pythreejs_viewer.py b/craftsman/utils/visualizers/pythreejs_viewer.py
deleted file mode 100755
index 70abbb7dce17361489f36eddca9c7e3089d4ec05..0000000000000000000000000000000000000000
--- a/craftsman/utils/visualizers/pythreejs_viewer.py
+++ /dev/null
@@ -1,537 +0,0 @@
-import numpy as np
-from ipywidgets import embed
-import pythreejs as p3s
-import uuid
-
-from .color_util import get_colors, gen_circle, gen_checkers
-
-
-EMBED_URL = "https://cdn.jsdelivr.net/npm/@jupyter-widgets/html-manager@1.0.1/dist/embed-amd.js"
-
-
-class PyThreeJSViewer(object):
-
- def __init__(self, settings, render_mode="WEBSITE"):
- self.render_mode = render_mode
- self.__update_settings(settings)
- self._light = p3s.DirectionalLight(color='white', position=[0, 0, 1], intensity=0.6)
- self._light2 = p3s.AmbientLight(intensity=0.5)
- self._cam = p3s.PerspectiveCamera(position=[0, 0, 1], lookAt=[0, 0, 0], fov=self.__s["fov"],
- aspect=self.__s["width"] / self.__s["height"], children=[self._light])
- self._orbit = p3s.OrbitControls(controlling=self._cam)
- self._scene = p3s.Scene(children=[self._cam, self._light2], background=self.__s["background"]) # "#4c4c80"
- self._renderer = p3s.Renderer(camera=self._cam, scene=self._scene, controls=[self._orbit],
- width=self.__s["width"], height=self.__s["height"],
- antialias=self.__s["antialias"])
-
- self.__objects = {}
- self.__cnt = 0
-
- def jupyter_mode(self):
- self.render_mode = "JUPYTER"
-
- def offline(self):
- self.render_mode = "OFFLINE"
-
- def website(self):
- self.render_mode = "WEBSITE"
-
- def __get_shading(self, shading):
- shad = {"flat": True, "wireframe": False, "wire_width": 0.03, "wire_color": "black",
- "side": 'DoubleSide', "colormap": "viridis", "normalize": [None, None],
- "bbox": False, "roughness": 0.5, "metalness": 0.25, "reflectivity": 1.0,
- "line_width": 1.0, "line_color": "black",
- "point_color": "red", "point_size": 0.01, "point_shape": "circle",
- "text_color": "red"
- }
- for k in shading:
- shad[k] = shading[k]
- return shad
-
- def __update_settings(self, settings={}):
- sett = {"width": 600, "height": 600, "antialias": True, "scale": 1.5, "background": "#ffffff",
- "fov": 30}
- for k in settings:
- sett[k] = settings[k]
- self.__s = sett
-
- def __add_object(self, obj, parent=None):
- if not parent: # Object is added to global scene and objects dict
- self.__objects[self.__cnt] = obj
- self.__cnt += 1
- self._scene.add(obj["mesh"])
- else: # Object is added to parent object and NOT to objects dict
- parent.add(obj["mesh"])
-
- self.__update_view()
-
- if self.render_mode == "JUPYTER":
- return self.__cnt - 1
- elif self.render_mode == "WEBSITE":
- return self
-
- def __add_line_geometry(self, lines, shading, obj=None):
- lines = lines.astype("float32", copy=False)
- mi = np.min(lines, axis=0)
- ma = np.max(lines, axis=0)
-
- geometry = p3s.LineSegmentsGeometry(positions=lines.reshape((-1, 2, 3)))
- material = p3s.LineMaterial(linewidth=shading["line_width"], color=shading["line_color"])
- # , vertexColors='VertexColors'),
- lines = p3s.LineSegments2(geometry=geometry, material=material) # type='LinePieces')
- line_obj = {"geometry": geometry, "mesh": lines, "material": material,
- "max": ma, "min": mi, "type": "Lines", "wireframe": None}
-
- if obj:
- return self.__add_object(line_obj, obj), line_obj
- else:
- return self.__add_object(line_obj)
-
- def __update_view(self):
- if len(self.__objects) == 0:
- return
- ma = np.zeros((len(self.__objects), 3))
- mi = np.zeros((len(self.__objects), 3))
- for r, obj in enumerate(self.__objects):
- ma[r] = self.__objects[obj]["max"]
- mi[r] = self.__objects[obj]["min"]
- ma = np.max(ma, axis=0)
- mi = np.min(mi, axis=0)
- diag = np.linalg.norm(ma - mi)
- mean = ((ma - mi) / 2 + mi).tolist()
- scale = self.__s["scale"] * (diag)
- self._orbit.target = mean
- self._cam.lookAt(mean)
- self._cam.position = [mean[0], mean[1], mean[2] + scale]
- self._light.position = [mean[0], mean[1], mean[2] + scale]
-
- self._orbit.exec_three_obj_method('update')
- self._cam.exec_three_obj_method('updateProjectionMatrix')
-
- def __get_bbox(self, v):
- m = np.min(v, axis=0)
- M = np.max(v, axis=0)
-
- # Corners of the bounding box
- v_box = np.array([[m[0], m[1], m[2]], [M[0], m[1], m[2]], [M[0], M[1], m[2]], [m[0], M[1], m[2]],
- [m[0], m[1], M[2]], [M[0], m[1], M[2]], [M[0], M[1], M[2]], [m[0], M[1], M[2]]])
-
- f_box = np.array([[0, 1], [1, 2], [2, 3], [3, 0], [4, 5], [5, 6], [6, 7], [7, 4],
- [0, 4], [1, 5], [2, 6], [7, 3]], dtype=np.uint32)
- return v_box, f_box
-
- def __get_colors(self, v, f, c, sh):
- coloring = "VertexColors"
- if type(c) == np.ndarray and c.size == 3: # Single color
- colors = np.ones_like(v)
- colors[:, 0] = c[0]
- colors[:, 1] = c[1]
- colors[:, 2] = c[2]
- # print("Single colors")
- elif type(c) == np.ndarray and len(c.shape) == 2 and c.shape[1] == 3: # Color values for
- if c.shape[0] == f.shape[0]: # faces
- colors = np.hstack([c, c, c]).reshape((-1, 3))
- coloring = "FaceColors"
- # print("Face color values")
- elif c.shape[0] == v.shape[0]: # vertices
- colors = c
- # print("Vertex color values")
- else: # Wrong size, fallback
- print("Invalid color array given! Supported are numpy arrays.", type(c))
- colors = np.ones_like(v)
- colors[:, 0] = 1.0
- colors[:, 1] = 0.874
- colors[:, 2] = 0.0
- elif type(c) == np.ndarray and c.size == f.shape[0]: # Function values for faces
- normalize = sh["normalize"][0] != None and sh["normalize"][1] != None
- cc = get_colors(c, sh["colormap"], normalize=normalize,
- vmin=sh["normalize"][0], vmax=sh["normalize"][1])
- # print(cc.shape)
- colors = np.hstack([cc, cc, cc]).reshape((-1, 3))
- coloring = "FaceColors"
- # print("Face function values")
- elif type(c) == np.ndarray and c.size == v.shape[0]: # Function values for vertices
- normalize = sh["normalize"][0] != None and sh["normalize"][1] != None
- colors = get_colors(c, sh["colormap"], normalize=normalize,
- vmin=sh["normalize"][0], vmax=sh["normalize"][1])
- # print("Vertex function values")
-
- else:
- colors = np.ones_like(v)
- # colors[:, 0] = 1.0
- # colors[:, 1] = 0.874
- # colors[:, 2] = 0.0
- colors[:, 0] = 1
- colors[:, 1] = 1
- colors[:, 2] = 1
-
- # No color
- if c is not None:
- print("Invalid color array given! Supported are numpy arrays.", type(c))
-
- return colors, coloring
-
- def __get_point_colors(self, v, c, sh):
- v_color = True
- if c is None: # No color given, use global color
- # conv = mpl.colors.ColorConverter()
- colors = sh["point_color"] # np.array(conv.to_rgb(sh["point_color"]))
- v_color = False
- elif isinstance(c, str): # No color given, use global color
- # conv = mpl.colors.ColorConverter()
- colors = c # np.array(conv.to_rgb(c))
- v_color = False
- elif type(c) == np.ndarray and len(c.shape) == 2 and c.shape[0] == v.shape[0] and c.shape[1] == 3:
- # Point color
- colors = c.astype("float32", copy=False)
-
- elif isinstance(c, np.ndarray) and len(c.shape) == 2 and c.shape[0] == v.shape[0] and c.shape[1] != 3:
- # Function values for vertices, but the colors are features
- c_norm = np.linalg.norm(c, ord=2, axis=-1)
- normalize = sh["normalize"][0] != None and sh["normalize"][1] != None
- colors = get_colors(c_norm, sh["colormap"], normalize=normalize,
- vmin=sh["normalize"][0], vmax=sh["normalize"][1])
- colors = colors.astype("float32", copy=False)
-
- elif type(c) == np.ndarray and c.size == v.shape[0]: # Function color
- normalize = sh["normalize"][0] != None and sh["normalize"][1] != None
- colors = get_colors(c, sh["colormap"], normalize=normalize,
- vmin=sh["normalize"][0], vmax=sh["normalize"][1])
- colors = colors.astype("float32", copy=False)
- # print("Vertex function values")
-
- else:
- print("Invalid color array given! Supported are numpy arrays.", type(c))
- colors = sh["point_color"]
- v_color = False
-
- return colors, v_color
-
- def add_mesh(self, v, f, c=None, uv=None, n=None, shading={}, texture_data=None, **kwargs):
- shading.update(kwargs)
- sh = self.__get_shading(shading)
- mesh_obj = {}
-
- # it is a tet
- if v.shape[1] == 3 and f.shape[1] == 4:
- f_tmp = np.ndarray([f.shape[0] * 4, 3], dtype=f.dtype)
- for i in range(f.shape[0]):
- f_tmp[i * 4 + 0] = np.array([f[i][1], f[i][0], f[i][2]])
- f_tmp[i * 4 + 1] = np.array([f[i][0], f[i][1], f[i][3]])
- f_tmp[i * 4 + 2] = np.array([f[i][1], f[i][2], f[i][3]])
- f_tmp[i * 4 + 3] = np.array([f[i][2], f[i][0], f[i][3]])
- f = f_tmp
-
- if v.shape[1] == 2:
- v = np.append(v, np.zeros([v.shape[0], 1]), 1)
-
- # Type adjustment vertices
- v = v.astype("float32", copy=False)
-
- # Color setup
- colors, coloring = self.__get_colors(v, f, c, sh)
-
- # Type adjustment faces and colors
- c = colors.astype("float32", copy=False)
-
- # Material and geometry setup
- ba_dict = {"color": p3s.BufferAttribute(c)}
- if coloring == "FaceColors":
- verts = np.zeros((f.shape[0] * 3, 3), dtype="float32")
- for ii in range(f.shape[0]):
- # print(ii*3, f[ii])
- verts[ii * 3] = v[f[ii, 0]]
- verts[ii * 3 + 1] = v[f[ii, 1]]
- verts[ii * 3 + 2] = v[f[ii, 2]]
- v = verts
- else:
- f = f.astype("uint32", copy=False).ravel()
- ba_dict["index"] = p3s.BufferAttribute(f, normalized=False)
-
- ba_dict["position"] = p3s.BufferAttribute(v, normalized=False)
-
- if uv is not None:
- uv = (uv - np.min(uv)) / (np.max(uv) - np.min(uv))
- if texture_data is None:
- texture_data = gen_checkers(20, 20)
- tex = p3s.DataTexture(data=texture_data, format="RGBFormat", type="FloatType")
- material = p3s.MeshStandardMaterial(map=tex, reflectivity=sh["reflectivity"], side=sh["side"],
- roughness=sh["roughness"], metalness=sh["metalness"],
- flatShading=sh["flat"],
- polygonOffset=True, polygonOffsetFactor=1, polygonOffsetUnits=5)
- ba_dict["uv"] = p3s.BufferAttribute(uv.astype("float32", copy=False))
- else:
- material = p3s.MeshStandardMaterial(vertexColors=coloring, reflectivity=sh["reflectivity"],
- side=sh["side"], roughness=sh["roughness"], metalness=sh["metalness"],
- flatShading=sh["flat"],
- polygonOffset=True, polygonOffsetFactor=1, polygonOffsetUnits=5)
-
- if type(n) != type(None) and coloring == "VertexColors": # TODO: properly handle normals for FaceColors as well
- ba_dict["normal"] = p3s.BufferAttribute(n.astype("float32", copy=False), normalized=True)
-
- geometry = p3s.BufferGeometry(attributes=ba_dict)
-
- if coloring == "VertexColors" and type(n) == type(None):
- geometry.exec_three_obj_method('computeVertexNormals')
- elif coloring == "FaceColors" and type(n) == type(None):
- geometry.exec_three_obj_method('computeFaceNormals')
-
- # Mesh setup
- mesh = p3s.Mesh(geometry=geometry, material=material)
-
- # Wireframe setup
- mesh_obj["wireframe"] = None
- if sh["wireframe"]:
- wf_geometry = p3s.WireframeGeometry(mesh.geometry) # WireframeGeometry
- wf_material = p3s.LineBasicMaterial(color=sh["wire_color"], linewidth=sh["wire_width"])
- wireframe = p3s.LineSegments(wf_geometry, wf_material)
- mesh.add(wireframe)
- mesh_obj["wireframe"] = wireframe
-
- # Bounding box setup
- if sh["bbox"]:
- v_box, f_box = self.__get_bbox(v)
- _, bbox = self.add_edges(v_box, f_box, sh, mesh)
- mesh_obj["bbox"] = [bbox, v_box, f_box]
-
- # Object setup
- mesh_obj["max"] = np.max(v, axis=0)
- mesh_obj["min"] = np.min(v, axis=0)
- mesh_obj["geometry"] = geometry
- mesh_obj["mesh"] = mesh
- mesh_obj["material"] = material
- mesh_obj["type"] = "Mesh"
- mesh_obj["shading"] = sh
- mesh_obj["coloring"] = coloring
- mesh_obj["arrays"] = [v, f, c] # TODO replays with proper storage or remove if not needed
-
- return self.__add_object(mesh_obj)
-
- def add_lines(self, beginning, ending, shading={}, obj=None, **kwargs):
- shading.update(kwargs)
- if len(beginning.shape) == 1:
- if len(beginning) == 2:
- beginning = np.array([[beginning[0], beginning[1], 0]])
- else:
- if beginning.shape[1] == 2:
- beginning = np.append(
- beginning, np.zeros([beginning.shape[0], 1]), 1)
- if len(ending.shape) == 1:
- if len(ending) == 2:
- ending = np.array([[ending[0], ending[1], 0]])
- else:
- if ending.shape[1] == 2:
- ending = np.append(
- ending, np.zeros([ending.shape[0], 1]), 1)
-
- sh = self.__get_shading(shading)
- lines = np.hstack([beginning, ending])
- lines = lines.reshape((-1, 3))
- return self.__add_line_geometry(lines, sh, obj)
-
- def add_edges(self, vertices, edges, shading={}, obj=None, **kwargs):
- shading.update(kwargs)
- if vertices.shape[1] == 2:
- vertices = np.append(
- vertices, np.zeros([vertices.shape[0], 1]), 1)
- sh = self.__get_shading(shading)
- lines = np.zeros((edges.size, 3))
- cnt = 0
- for e in edges:
- lines[cnt, :] = vertices[e[0]]
- lines[cnt + 1, :] = vertices[e[1]]
- cnt += 2
- return self.__add_line_geometry(lines, sh, obj)
-
- def add_points(self, points, c=None, shading={}, obj=None, **kwargs):
- shading.update(kwargs)
- if len(points.shape) == 1:
- if len(points) == 2:
- points = np.array([[points[0], points[1], 0]])
- else:
- if points.shape[1] == 2:
- points = np.append(
- points, np.zeros([points.shape[0], 1]), 1)
- sh = self.__get_shading(shading)
- points = points.astype("float32", copy=False)
- mi = np.min(points, axis=0)
- ma = np.max(points, axis=0)
-
- g_attributes = {"position": p3s.BufferAttribute(points, normalized=False)}
- m_attributes = {"size": sh["point_size"]}
-
- if sh["point_shape"] == "circle": # Plot circles
- tex = p3s.DataTexture(data=gen_circle(16, 16), format="RGBAFormat", type="FloatType")
- m_attributes["map"] = tex
- m_attributes["alphaTest"] = 0.5
- m_attributes["transparency"] = True
- else: # Plot squares
- pass
-
- colors, v_colors = self.__get_point_colors(points, c, sh)
- if v_colors: # Colors per point
- m_attributes["vertexColors"] = 'VertexColors'
- g_attributes["color"] = p3s.BufferAttribute(colors, normalized=False)
-
- else: # Colors for all points
- m_attributes["color"] = colors
-
- material = p3s.PointsMaterial(**m_attributes)
- geometry = p3s.BufferGeometry(attributes=g_attributes)
- points = p3s.Points(geometry=geometry, material=material)
- point_obj = {"geometry": geometry, "mesh": points, "material": material,
- "max": ma, "min": mi, "type": "Points", "wireframe": None}
-
- if obj:
- return self.__add_object(point_obj, obj), point_obj
- else:
- return self.__add_object(point_obj)
-
- def remove_object(self, obj_id):
- if obj_id not in self.__objects:
- print("Invalid object id. Valid ids are: ", list(self.__objects.keys()))
- return
- self._scene.remove(self.__objects[obj_id]["mesh"])
- del self.__objects[obj_id]
- self.__update_view()
-
- def reset(self):
- for obj_id in list(self.__objects.keys()).copy():
- self._scene.remove(self.__objects[obj_id]["mesh"])
- del self.__objects[obj_id]
- self.__update_view()
-
- def update_object(self, oid=0, vertices=None, colors=None, faces=None):
- obj = self.__objects[oid]
- if type(vertices) != type(None):
- if obj["coloring"] == "FaceColors":
- f = obj["arrays"][1]
- verts = np.zeros((f.shape[0] * 3, 3), dtype="float32")
- for ii in range(f.shape[0]):
- # print(ii*3, f[ii])
- verts[ii * 3] = vertices[f[ii, 0]]
- verts[ii * 3 + 1] = vertices[f[ii, 1]]
- verts[ii * 3 + 2] = vertices[f[ii, 2]]
- v = verts
-
- else:
- v = vertices.astype("float32", copy=False)
- obj["geometry"].attributes["position"].array = v
- # self.wireframe.attributes["position"].array = v # Wireframe updates?
- obj["geometry"].attributes["position"].needsUpdate = True
- # obj["geometry"].exec_three_obj_method('computeVertexNormals')
- if type(colors) != type(None):
- colors, coloring = self.__get_colors(obj["arrays"][0], obj["arrays"][1], colors, obj["shading"])
- colors = colors.astype("float32", copy=False)
- obj["geometry"].attributes["color"].array = colors
- obj["geometry"].attributes["color"].needsUpdate = True
- if type(faces) != type(None):
- if obj["coloring"] == "FaceColors":
- print("Face updates are currently only possible in vertex color mode.")
- return
- f = faces.astype("uint32", copy=False).ravel()
- print(obj["geometry"].attributes)
- obj["geometry"].attributes["index"].array = f
- # self.wireframe.attributes["position"].array = v # Wireframe updates?
- obj["geometry"].attributes["index"].needsUpdate = True
- # obj["geometry"].exec_three_obj_method('computeVertexNormals')
- # self.mesh.geometry.verticesNeedUpdate = True
- # self.mesh.geometry.elementsNeedUpdate = True
- # self.update()
- if self.render_mode == "WEBSITE":
- return self
-
- # def update(self):
- # self.mesh.exec_three_obj_method('update')
- # self.orbit.exec_three_obj_method('update')
- # self.cam.exec_three_obj_method('updateProjectionMatrix')
- # self.scene.exec_three_obj_method('update')
-
- def add_text(self, text, shading={}, **kwargs):
- shading.update(kwargs)
- sh = self.__get_shading(shading)
- tt = p3s.TextTexture(string=text, color=sh["text_color"])
- sm = p3s.SpriteMaterial(map=tt)
- text = p3s.Sprite(material=sm, scaleToTexture=True)
- self._scene.add(text)
-
- # def add_widget(self, widget, callback):
- # self.widgets.append(widget)
- # widget.observe(callback, names='value')
-
- # def add_dropdown(self, options, default, desc, cb):
- # widget = widgets.Dropdown(options=options, value=default, description=desc)
- # self.__widgets.append(widget)
- # widget.observe(cb, names="value")
- # display(widget)
-
- # def add_button(self, text, cb):
- # button = widgets.Button(description=text)
- # self.__widgets.append(button)
- # button.on_click(cb)
- # display(button)
-
- def to_html(self, imports=True, html_frame=True):
- # Bake positions (fixes centering bug in offline rendering)
- if len(self.__objects) == 0:
- return
- ma = np.zeros((len(self.__objects), 3))
- mi = np.zeros((len(self.__objects), 3))
- for r, obj in enumerate(self.__objects):
- ma[r] = self.__objects[obj]["max"]
- mi[r] = self.__objects[obj]["min"]
- ma = np.max(ma, axis=0)
- mi = np.min(mi, axis=0)
- diag = np.linalg.norm(ma - mi)
- mean = (ma - mi) / 2 + mi
- for r, obj in enumerate(self.__objects):
- v = self.__objects[obj]["geometry"].attributes["position"].array
- v -= mean
- v += np.array([0.0, .9, 0.0]) #! to move the obj to the center of window
-
- scale = self.__s["scale"] * (diag)
- self._orbit.target = [0.0, 0.0, 0.0]
- self._cam.lookAt([0.0, 0.0, 0.0])
- # self._cam.position = [0.0, 0.0, scale]
- self._cam.position = [0.0, 0.5, scale * 1.3] #! show four complete meshes in the window
- self._light.position = [0.0, 0.0, scale]
-
- state = embed.dependency_state(self._renderer)
-
- # Somehow these entries are missing when the state is exported in python.
- # Exporting from the GUI works, so we are inserting the missing entries.
- for k in state:
- if state[k]["model_name"] == "OrbitControlsModel":
- state[k]["state"]["maxAzimuthAngle"] = "inf"
- state[k]["state"]["maxDistance"] = "inf"
- state[k]["state"]["maxZoom"] = "inf"
- state[k]["state"]["minAzimuthAngle"] = "-inf"
-
- tpl = embed.load_requirejs_template
- if not imports:
- embed.load_requirejs_template = ""
-
- s = embed.embed_snippet(self._renderer, state=state, embed_url=EMBED_URL)
- # s = embed.embed_snippet(self.__w, state=state)
- embed.load_requirejs_template = tpl
-
- if html_frame:
- s = "\n\n" + s + "\n\n"
-
- # Revert changes
- for r, obj in enumerate(self.__objects):
- v = self.__objects[obj]["geometry"].attributes["position"].array
- v += mean
- self.__update_view()
-
- return s
-
- def save(self, filename=""):
- if filename == "":
- uid = str(uuid.uuid4()) + ".html"
- else:
- filename = filename.replace(".html", "")
- uid = filename + '.html'
- with open(uid, "w") as f:
- f.write(self.to_html())
- print("Plot saved to file %s." % uid)
diff --git a/examples/bug.png b/examples/bug.png
new file mode 100644
index 0000000000000000000000000000000000000000..bd730216d50e830c4010f08a30442f626268a3ac
Binary files /dev/null and b/examples/bug.png differ
diff --git a/examples/dragon.png b/examples/dragon.png
new file mode 100644
index 0000000000000000000000000000000000000000..1ed16896722535aedbdaf629dd0068cd168bd597
Binary files /dev/null and b/examples/dragon.png differ
diff --git a/examples/mask.png b/examples/mask.png
new file mode 100644
index 0000000000000000000000000000000000000000..6ab444c7a2850ac645567ff6f93ed8149d9cbb45
Binary files /dev/null and b/examples/mask.png differ
diff --git a/examples/monster.png b/examples/monster.png
new file mode 100644
index 0000000000000000000000000000000000000000..1cfe5eb404319480d5f17390fd4c76baa890b300
Binary files /dev/null and b/examples/monster.png differ
diff --git a/examples/werewolf.png b/examples/werewolf.png
new file mode 100644
index 0000000000000000000000000000000000000000..f20d321902a1c50f642533dbe90ee1aa83039ccd
Binary files /dev/null and b/examples/werewolf.png differ
diff --git a/gradio_app.py b/gradio_app.py
old mode 100644
new mode 100755
index c5944f7bae5525247a2b26d9934fc4c31912e47b..3e2e0140e0d81aac04fd03a5c5af373a1fd9bf39
--- a/gradio_app.py
+++ b/gradio_app.py
@@ -9,26 +9,32 @@ import importlib
import numpy as np
from omegaconf import OmegaConf
from huggingface_hub import hf_hub_download
+from diffusers import DiffusionPipeline
+import PIL
+from PIL import Image
from collections import OrderedDict
import trimesh
+import rembg
import gradio as gr
from typing import Any
-from einops import rearrange
proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(os.path.join(proj_dir))
import tempfile
-from apps.utils import *
+import craftsman
+from craftsman.utils.config import ExperimentConfig, load_config
_TITLE = '''CraftsMan: High-fidelity Mesh Generation with 3D Native Generation and Interactive Geometry Refiner'''
_DESCRIPTION = '''
+Select or upload a image, then just click 'Generate'.
-By mimicking the artist/craftsman modeling workflow, we propose CraftsMan (aka 匠心) which uses 3D Latent Set Diffusion Model that directly generates coarse meshes, +By mimicking the artist/craftsman modeling workflow, we propose CraftsMan (aka 匠心) that uses 3D Latent Set Diffusion Model that directly generate coarse meshes, then a multi-view normal enhanced image generation model is used to refine the mesh. We provide the coarse 3D diffusion part here.
@@ -36,8 +42,6 @@ If you found CraftsMan is helpful, please help to ⭐ the
*If you have your own multi-view images, you can directly upload it. -Tutorial -使用教程
'''
_CITE_ = r"""
@@ -59,153 +63,101 @@ CraftsMan is under [AGPL-3.0](https://www.gnu.org/licenses/agpl-3.0.en.html), so
📧 **Contact**
If you have any questions, feel free to open a discussion or contact us at weiyuli.cn@gmail.com.
"""
-from apps.third_party.CRM.pipelines import TwoStagePipeline
-from apps.third_party.LGM.pipeline_mvdream import MVDreamPipeline
-from apps.third_party.Era3D.pipelines.pipeline_mvdiffusion_unclip import StableUnCLIPImg2ImgPipeline
-from apps.third_party.Era3D.data.single_image_dataset import SingleImageDataset
-
-import re
-import os
-import stat
-
-RD, WD, XD = 4, 2, 1
-BNS = [RD, WD, XD]
-MDS = [
- [stat.S_IRUSR, stat.S_IRGRP, stat.S_IROTH],
- [stat.S_IWUSR, stat.S_IWGRP, stat.S_IWOTH],
- [stat.S_IXUSR, stat.S_IXGRP, stat.S_IXOTH]
-]
-
-def chmod(path, mode):
- if isinstance(mode, int):
- mode = str(mode)
- if not re.match("^[0-7]{1,3}$", mode):
- raise Exception("mode does not conform to ^[0-7]{1,3}$ pattern")
- mode = "{0:0>3}".format(mode)
- mode_num = 0
- for midx, m in enumerate(mode):
- for bnidx, bn in enumerate(BNS):
- if (int(m) & bn) > 0:
- mode_num += MDS[bnidx][midx]
- os.chmod(path, mode_num)
-
-chmod(f"{parent_dir}/apps/third_party/InstantMeshes", "777")
-device = None
model = None
cached_dir = None
-generator = None
-
-sys.path.append(f"apps/third_party/CRM")
-crm_pipeline = None
-
-sys.path.append(f"apps/third_party/LGM")
-imgaedream_pipeline = None
-
-sys.path.append(f"apps/third_party/Era3D")
-era3d_pipeline = None
-
-@spaces.GPU(duration=120)
-def gen_mvimg(
- mvimg_model, image, seed, guidance_scale, step, text, neg_text, elevation, backgroud_color
-):
- global device
- if seed == 0:
- seed = np.random.randint(1, 65535)
- global generator
- generator = torch.Generator(device)
- generator.manual_seed(seed)
- if mvimg_model == "CRM":
- global crm_pipeline
- crm_pipeline.set_seed(seed)
- background = Image.new("RGBA", image.size, (127, 127, 127))
- image = Image.alpha_composite(background, image)
- mv_imgs = crm_pipeline(
- image,
- scale=guidance_scale,
- step=step
- )["stage1_images"]
- return mv_imgs[5], mv_imgs[3], mv_imgs[2], mv_imgs[0]
-
- elif mvimg_model == "ImageDream":
- global imagedream_pipeline
- background = Image.new("RGBA", image.size, backgroud_color)
- image = Image.alpha_composite(background, image)
- image = np.array(image).astype(np.float32) / 255.0
- image = image[..., :3] * image[..., 3:4] + (1 - image[..., 3:4])
- mv_imgs = imagedream_pipeline(
- text,
- image,
- negative_prompt=neg_text,
- guidance_scale=guidance_scale,
- num_inference_steps=step,
- elevation=elevation,
- generator=generator,
- )
- return mv_imgs[1], mv_imgs[2], mv_imgs[3], mv_imgs[0]
-
- elif mvimg_model == "Era3D":
- global era3d_pipeline
- era3d_pipeline.to(device)
- era3d_pipeline.unet.enable_xformers_memory_efficient_attention()
- era3d_pipeline.set_progress_bar_config(disable=True)
+generator = None
- crop_size = 420
- batch = SingleImageDataset(root_dir='', num_views=6, img_wh=[512, 512], bg_color='white',
- crop_size=crop_size, single_image=image, prompt_embeds_path='apps/third_party/Era3D/data/fixed_prompt_embeds_6view')[0]
- imgs_in = torch.cat([batch['imgs_in']]*2, dim=0)
- imgs_in = rearrange(imgs_in, "B Nv C H W -> (B Nv) C H W")# (B*Nv, 3, H, W)
-
- normal_prompt_embeddings, clr_prompt_embeddings = batch['normal_prompt_embeddings'], batch['color_prompt_embeddings']
- prompt_embeddings = torch.cat([normal_prompt_embeddings, clr_prompt_embeddings], dim=0)
- prompt_embeddings = rearrange(prompt_embeddings, "B Nv N C -> (B Nv) N C")
-
- imgs_in = imgs_in.to(dtype=torch.float16)
- prompt_embeddings = prompt_embeddings.to(dtype=torch.float16)
+def check_input_image(input_image):
+ if input_image is None:
+ raise gr.Error("No image uploaded!")
+
+class RMBG(object):
+ def __init__(self):
+ pass
+
+ def rmbg_rembg(self, input_image, background_color):
+ def _rembg_remove(
+ image: PIL.Image.Image,
+ rembg_session = None,
+ force: bool = False,
+ **rembg_kwargs,
+ ) -> PIL.Image.Image:
+ do_remove = True
+ if image.mode == "RGBA" and image.getextrema()[3][0] < 255:
+ # explain why current do not rm bg
+ print("alhpa channl not enpty, skip remove background, using alpha channel as mask")
+ background = Image.new("RGBA", image.size, background_color)
+ image = Image.alpha_composite(background, image)
+ do_remove = False
+ do_remove = do_remove or force
+ if do_remove:
+ image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
+
+ # calculate the min bbox of the image
+ alpha = image.split()[-1]
+ image = image.crop(alpha.getbbox())
+
+ return image
+ return _rembg_remove(input_image, None, force_remove=True)
+
+ def run(self, rm_type, image, foreground_ratio, background_choice, background_color=(0, 0, 0, 0)):
+ if "Original" in background_choice:
+ return image
+ else:
+ if background_choice == "Alpha as mask":
+ alpha = image.split()[-1]
+ image = image.crop(alpha.getbbox())
+
+ elif "Remove" in background_choice:
+ if rm_type.upper() == "REMBG":
+ image = self.rmbg_rembg(image, background_color=background_color)
+ else:
+ return -1
- mv_imgs = era3d_pipeline(
- imgs_in,
- None,
- prompt_embeds=prompt_embeddings,
- generator=generator,
- guidance_scale=guidance_scale,
- num_inference_steps=step,
- num_images_per_prompt=1,
- **{'eta': 1.0}
- ).images
- return mv_imgs[6], mv_imgs[8], mv_imgs[9], mv_imgs[10]
-
-@spaces.GPU
-def image2mesh(view_front: np.ndarray,
- view_right: np.ndarray,
- view_back: np.ndarray,
- view_left: np.ndarray,
+ # Calculate the new size after rescaling
+ new_size = tuple(int(dim * foreground_ratio) for dim in image.size)
+ # Resize the image while maintaining the aspect ratio
+ resized_image = image.resize(new_size)
+ # Create a new image with the original size and white background
+ padded_image = PIL.Image.new("RGBA", image.size, (0, 0, 0, 0))
+ paste_position = ((image.width - resized_image.width) // 2, (image.height - resized_image.height) // 2)
+ padded_image.paste(resized_image, paste_position)
+
+ # expand image to 1:1
+ width, height = padded_image.size
+ if width == height:
+ return padded_image
+ new_size = (max(width, height), max(width, height))
+ image = PIL.Image.new("RGBA", new_size, (0, 0, 0, 0))
+ paste_position = ((new_size[0] - width) // 2, (new_size[1] - height) // 2)
+ image.paste(padded_image, paste_position)
+ return image
+
+# @spaces.GPU
+def image2mesh(image: Any,
more: bool = False,
scheluder_name: str ="DDIMScheduler",
guidance_scale: int = 7.5,
- steps: int = 50,
+ steps: int = 30,
seed: int = 4,
+ target_face_count: int = 2000,
octree_depth: int = 7):
sample_inputs = {
- "mvimages": [[
- Image.fromarray(view_front),
- Image.fromarray(view_right),
- Image.fromarray(view_back),
- Image.fromarray(view_left)
- ]]
+ "image": [
+ image
+ ]
}
global model
latents = model.sample(
sample_inputs,
sample_times=1,
- guidance_scale=guidance_scale,
- return_intermediates=False,
steps=steps,
+ guidance_scale=guidance_scale,
seed=seed
-
)[0]
# decode the latents to mesh
@@ -224,75 +176,42 @@ def image2mesh(view_front: np.ndarray,
if 'Remesh' in more:
remeshed_filepath = tempfile.NamedTemporaryFile(suffix=f"_remeshed.obj", delete=False).name
print("Remeshing with Instant Meshes...")
- # target_face_count = int(len(mesh.faces)/10)
- target_face_count = 2000
command = f"{proj_dir}/apps/third_party/InstantMeshes {filepath} -f {target_face_count} -o {remeshed_filepath}"
os.system(command)
- del filepath
filepath = remeshed_filepath
- # filepath = filepath.replace('.obj', '_remeshed.obj')
return filepath
if __name__=="__main__":
parser = argparse.ArgumentParser()
- # parser.add_argument("--model_path", type=str, required=True, help="Path to the object file",)
- parser.add_argument("--cached_dir", type=str, default="./gradio_cached_dir")
+ parser.add_argument("--model_path", type=str, default="./ckpts/craftsman-v1-5", help="Path to the object file",)
+ parser.add_argument("--cached_dir", type=str, default="")
parser.add_argument("--device", type=int, default=0)
args = parser.parse_args()
+
cached_dir = args.cached_dir
- os.makedirs(args.cached_dir, exist_ok=True)
+ if cached_dir != "":
+ os.makedirs(args.cached_dir, exist_ok=True)
device = torch.device(f"cuda:{args.device}" if torch.cuda.is_available() else "cpu")
print(f"using device: {device}")
- # for multi-view images generation
+ # for input image
background_choice = OrderedDict({
"Alpha as Mask": "Alpha as Mask",
"Auto Remove Background": "Auto Remove Background",
"Original Image": "Original Image",
})
- mvimg_model_config_list = [
- "Era3D",
- "CRM",
- "ImageDream"
- ]
- if "Era3D" in mvimg_model_config_list:
- # cfg = load_config("apps/third_party/Era3D/configs/test_unclip-512-6view.yaml")
- # schema = OmegaConf.structured(TestConfig)
- # cfg = OmegaConf.merge(schema, cfg)
- era3d_pipeline = StableUnCLIPImg2ImgPipeline.from_pretrained(
- 'pengHTYX/MacLab-Era3D-512-6view',
- dtype=torch.float16,
- )
- # enable xformers
- # era3d_pipeline.unet.enable_xformers_memory_efficient_attention()
- # era3d_pipeline.to(device)
- if "CRM" in mvimg_model_config_list:
- stage1_config = OmegaConf.load(f"apps/third_party/CRM/configs/nf7_v3_SNR_rd_size_stroke.yaml").config
- stage1_sampler_config = stage1_config.sampler
- stage1_model_config = stage1_config.models
- stage1_model_config.resume = hf_hub_download(repo_id="Zhengyi/CRM", filename="pixel-diffusion.pth", repo_type="model")
- stage1_model_config.config = f"apps/third_party/CRM/" + stage1_model_config.config
- crm_pipeline = TwoStagePipeline(
- stage1_model_config,
- stage1_sampler_config,
- device=device,
- dtype=torch.float16
- )
- if "ImageDream" in mvimg_model_config_list:
- imagedream_pipeline = MVDreamPipeline.from_pretrained(
- "ashawkey/imagedream-ipmv-diffusers", # remote weights
- torch_dtype=torch.float16,
- trust_remote_code=True,
- )
+ generator = torch.Generator(device)
# for 3D latent set diffusion
- ckpt_path = hf_hub_download(repo_id="wyysf/CraftsMan", filename="image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-aligned-vae/model.ckpt", repo_type="model")
- config_path = hf_hub_download(repo_id="wyysf/CraftsMan", filename="image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-aligned-vae/config.yaml", repo_type="model")
- # ckpt_path = hf_hub_download(repo_id="wyysf/CraftsMan", filename="image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6/model-300k.ckpt", repo_type="model")
- # config_path = hf_hub_download(repo_id="wyysf/CraftsMan", filename="image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6/config.yaml", repo_type="model")
+ if args.model_path == "":
+ ckpt_path = hf_hub_download(repo_id="craftsman3d/craftsman-v1-5", filename="model.ckpt", repo_type="model")
+ config_path = hf_hub_download(repo_id="craftsman3d/craftsman-v1-5", filename="config.yaml", repo_type="model")
+ else:
+ ckpt_path = os.path.join(args.model_path, "model.ckpt")
+ config_path = os.path.join(args.model_path, "config.yaml")
scheluder_dict = OrderedDict({
"DDIMScheduler": 'diffusers.schedulers.DDIMScheduler',
# "DPMSolverMultistepScheduler": 'diffusers.schedulers.DPMSolverMultistepScheduler', # not support yet
@@ -330,22 +249,12 @@ if __name__=="__main__":
gr.Markdown('''Try a different seed and MV Model for better results. Good Luck :)''')
with gr.Row():
seed = gr.Number(0, label='Seed', show_label=True)
- mvimg_model = gr.Dropdown(value="CRM", label="MV Image Model", choices=list(mvimg_model_config_list))
more = gr.CheckboxGroup(["Remesh"], label="More", show_label=False)
-
- with gr.Row():
- # input prompt
- text = gr.Textbox(label="Prompt (Opt.)", info="only works for ImageDream")
-
- with gr.Accordion('Advanced options', open=False):
- # negative prompt
- neg_text = gr.Textbox(label="Negative Prompt", value='ugly, blurry, pixelated obscure, unnatural colors, poor lighting, dull, unclear, cropped, lowres, low quality, artifacts, duplicate')
- # elevation
- elevation = gr.Slider(label="elevation", minimum=-90, maximum=90, step=1, value=0)
+ target_face_count = gr.Number(2000, label='Target Face Count', show_label=True)
with gr.Row():
gr.Examples(
- examples=[os.path.join("./apps/examples", i) for i in os.listdir("./apps/examples")],
+ examples=[os.path.join("./assets/examples", i) for i in os.listdir("./assets/examples")],
inputs=[image_input],
examples_per_page=8
)
@@ -357,43 +266,17 @@ if __name__=="__main__":
camera_position=(90.0, 90.0, 3.5),
interactive=False,
)
- # with gr.Row():
- # gr.Markdown('''*please note that the model is fliped due to the gradio viewer, please download the obj file and you will get the correct orientation.''')
-
with gr.Row():
- view_front = gr.Image(label="Front", interactive=True, show_label=True)
- view_right = gr.Image(label="Right", interactive=True, show_label=True)
- view_back = gr.Image(label="Back", interactive=True, show_label=True)
- view_left = gr.Image(label="Left", interactive=True, show_label=True)
+ gr.Markdown('''*please note that the model is fliped due to the gradio viewer, please download the obj file and you will get the correct orientation.''')
with gr.Accordion('Advanced options', open=False):
- with gr.Row(equal_height=True):
- run_mv_btn = gr.Button('Only Generate 2D', interactive=True)
- run_3d_btn = gr.Button('Only Generate 3D', interactive=True)
-
- with gr.Accordion('Advanced options (2D)', open=False):
- with gr.Row():
- foreground_ratio = gr.Slider(
- label="Foreground Ratio",
- minimum=0.5,
- maximum=1.0,
- value=1.0,
- step=0.05,
- )
-
with gr.Row():
background_choice = gr.Dropdown(label="Backgroud Choice", value="Auto Remove Background",choices=list(background_choice.keys()))
rmbg_type = gr.Dropdown(label="Backgroud Remove Type", value="rembg",choices=['sam', "rembg"])
- backgroud_color = gr.ColorPicker(label="Background Color", value="#FFFFFF", interactive=True)
- # backgroud_color = gr.ColorPicker(label="Background Color", value="#7F7F7F", interactive=True)
-
- with gr.Row():
- mvimg_guidance_scale = gr.Number(value=3.0, minimum=1, maximum=10, label="2D Guidance Scale")
- mvimg_steps = gr.Number(value=30, minimum=20, maximum=100, label="2D Sample Steps")
-
- with gr.Accordion('Advanced options (3D)', open=False):
+ foreground_ratio = gr.Slider(label="Foreground Ratio", value=1.0, minimum=0.5, maximum=1.0, step=0.01)
+
with gr.Row():
- guidance_scale = gr.Number(label="3D Guidance Scale", value=3.0, minimum=1.0, maximum=10.0)
+ guidance_scale = gr.Number(label="3D Guidance Scale", value=5.0, minimum=3.0, maximum=10.0)
steps = gr.Number(value=50, minimum=20, maximum=100, label="3D Sample Steps")
with gr.Row():
@@ -403,31 +286,28 @@ if __name__=="__main__":
gr.Markdown(_CITE_)
outputs = [output_model_obj]
- rmbg = RMBG(device)
+ rmbg = RMBG()
- model = load_model(ckpt_path, config_path, device)
+ # model = load_model(ckpt_path, config_path, device)
+ cfg = load_config(config_path)
+ model = craftsman.find(cfg.system_type)(cfg.system)
+ print(f"Restoring states from the checkpoint path at {ckpt_path} with config {cfg}")
+ ckpt = torch.load(ckpt_path, map_location=torch.device('cpu'))
+ model.load_state_dict(
+ ckpt["state_dict"] if "state_dict" in ckpt else ckpt,
+ )
+ model = model.to(device).eval()
+
run_btn.click(fn=check_input_image, inputs=[image_input]
).success(
fn=rmbg.run,
- inputs=[rmbg_type, image_input, foreground_ratio, background_choice, backgroud_color],
+ inputs=[rmbg_type, image_input, foreground_ratio, background_choice],
outputs=[image_input]
- ).success(
- fn=gen_mvimg,
- inputs=[mvimg_model, image_input, seed, mvimg_guidance_scale, mvimg_steps, text, neg_text, elevation, backgroud_color],
- outputs=[view_front, view_right, view_back, view_left]
).success(
fn=image2mesh,
- inputs=[view_front, view_right, view_back, view_left, more, scheduler, guidance_scale, steps, seed, octree_depth],
+ inputs=[image_input, more, scheduler, guidance_scale, steps, seed, target_face_count, octree_depth],
outputs=outputs,
api_name="generate_img2obj")
- run_mv_btn.click(fn=gen_mvimg,
- inputs=[mvimg_model, image_input, seed, mvimg_guidance_scale, mvimg_steps, text, neg_text, elevation, backgroud_color],
- outputs=[view_front, view_right, view_back, view_left]
- )
- run_3d_btn.click(fn=image2mesh,
- inputs=[view_front, view_right, view_back, view_left, more, scheduler, guidance_scale, steps, seed, octree_depth],
- outputs=outputs,
- api_name="generate_img2obj")
demo.queue().launch(share=True, allowed_paths=[args.cached_dir])
\ No newline at end of file
diff --git a/launch.py b/launch.py
deleted file mode 100644
index 854425983580adb24a56bc56071a9003e087f47d..0000000000000000000000000000000000000000
--- a/launch.py
+++ /dev/null
@@ -1,299 +0,0 @@
-import argparse
-import contextlib
-import importlib
-import logging
-import os
-import sys
-import time
-import traceback
-import pytorch_lightning as pl
-import torch
-from pytorch_lightning import Trainer
-from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
-from pytorch_lightning.loggers import CSVLogger, TensorBoardLogger
-from pytorch_lightning.utilities.rank_zero import rank_zero_only
-import craftsman
-from craftsman.systems.base import BaseSystem
-from craftsman.utils.callbacks import (
- CodeSnapshotCallback,
- ConfigSnapshotCallback,
- CustomProgressBar,
- ProgressCallback,
-)
-from craftsman.utils.config import ExperimentConfig, load_config
-from craftsman.utils.misc import get_rank
-from craftsman.utils.typing import Optional
-class ColoredFilter(logging.Filter):
- """
- A logging filter to add color to certain log levels.
- """
-
- RESET = "\033[0m"
- RED = "\033[31m"
- GREEN = "\033[32m"
- YELLOW = "\033[33m"
- BLUE = "\033[34m"
- MAGENTA = "\033[35m"
- CYAN = "\033[36m"
-
- COLORS = {
- "WARNING": YELLOW,
- "INFO": GREEN,
- "DEBUG": BLUE,
- "CRITICAL": MAGENTA,
- "ERROR": RED,
- }
-
- RESET = "\x1b[0m"
-
- def __init__(self):
- super().__init__()
-
- def filter(self, record):
- if record.levelname in self.COLORS:
- color_start = self.COLORS[record.levelname]
- record.levelname = f"{color_start}[{record.levelname}]"
- record.msg = f"{record.msg}{self.RESET}"
- return True
-
-
-def load_custom_module(module_path):
- module_name = os.path.basename(module_path)
- if os.path.isfile(module_path):
- sp = os.path.splitext(module_path)
- module_name = sp[0]
- try:
- if os.path.isfile(module_path):
- module_spec = importlib.util.spec_from_file_location(
- module_name, module_path
- )
- else:
- module_spec = importlib.util.spec_from_file_location(
- module_name, os.path.join(module_path, "__init__.py")
- )
-
- module = importlib.util.module_from_spec(module_spec)
- sys.modules[module_name] = module
- module_spec.loader.exec_module(module)
- return True
- except Exception as e:
- print(traceback.format_exc())
- print(f"Cannot import {module_path} module for custom nodes:", e)
- return False
-
-
-def load_custom_modules():
- node_paths = ["custom"]
- node_import_times = []
- if not os.path.exists("node_paths"):
- return
- for custom_node_path in node_paths:
- possible_modules = os.listdir(custom_node_path)
- if "__pycache__" in possible_modules:
- possible_modules.remove("__pycache__")
-
- for possible_module in possible_modules:
- module_path = os.path.join(custom_node_path, possible_module)
- if (
- os.path.isfile(module_path)
- and os.path.splitext(module_path)[1] != ".py"
- ):
- continue
- if module_path.endswith(".disabled"):
- continue
- time_before = time.perf_counter()
- success = load_custom_module(module_path)
- node_import_times.append(
- (time.perf_counter() - time_before, module_path, success)
- )
-
- if len(node_import_times) > 0:
- print("\nImport times for custom modules:")
- for n in sorted(node_import_times):
- if n[2]:
- import_message = ""
- else:
- import_message = " (IMPORT FAILED)"
- print("{:6.1f} seconds{}:".format(n[0], import_message), n[1])
- print()
-
-
-def main(args, extras) -> None:
- # set CUDA_VISIBLE_DEVICES if needed, then import pytorch-lightning
- os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
- env_gpus_str = os.environ.get("CUDA_VISIBLE_DEVICES", None)
- env_gpus = list(env_gpus_str.split(",")) if env_gpus_str else []
- selected_gpus = [0]
- torch.set_float32_matmul_precision("high")
-
- # Always rely on CUDA_VISIBLE_DEVICES if specific GPU ID(s) are specified.
- # As far as Pytorch Lightning is concerned, we always use all available GPUs
- # (possibly filtered by CUDA_VISIBLE_DEVICES).
- devices = -1
- if len(env_gpus) > 0:
- n_gpus = len(env_gpus)
- else:
- selected_gpus = list(args.gpu.split(","))
- n_gpus = len(selected_gpus)
- print(f"Using {n_gpus} GPUs: {selected_gpus}")
- os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
-
- if args.typecheck:
- from jaxtyping import install_import_hook
-
- install_import_hook("craftsman", "typeguard.typechecked")
-
- logger = logging.getLogger("pytorch_lightning")
- if args.verbose:
- logger.setLevel(logging.DEBUG)
-
- for handler in logger.handlers:
- if handler.stream == sys.stderr: # type: ignore
- if not args.gradio:
- handler.setFormatter(logging.Formatter("%(levelname)s %(message)s"))
- handler.addFilter(ColoredFilter())
- else:
- handler.setFormatter(logging.Formatter("[%(levelname)s] %(message)s"))
-
- load_custom_modules()
-
- # parse YAML config to OmegaConf
- cfg: ExperimentConfig
- cfg = load_config(args.config, cli_args=extras, n_gpus=n_gpus)
-
- # set a different seed for each device
- pl.seed_everything(cfg.seed + get_rank(), workers=True)
-
- dm = craftsman.find(cfg.data_type)(cfg.data)
- system: BaseSystem = craftsman.find(cfg.system_type)(
- cfg.system, resumed=cfg.resume is not None
- )
- system.set_save_dir(os.path.join(cfg.trial_dir, "save"))
-
- if args.gradio:
- fh = logging.FileHandler(os.path.join(cfg.trial_dir, "logs"))
- fh.setLevel(logging.INFO)
- if args.verbose:
- fh.setLevel(logging.DEBUG)
- fh.setFormatter(logging.Formatter("[%(levelname)s] %(message)s"))
- logger.addHandler(fh)
-
- callbacks = []
- if args.train:
- callbacks += [
- ModelCheckpoint(
- dirpath=os.path.join(cfg.trial_dir, "ckpts"), **cfg.checkpoint
- ),
- LearningRateMonitor(logging_interval="step"),
- CodeSnapshotCallback(
- os.path.join(cfg.trial_dir, "code"), use_version=False
- ),
- ConfigSnapshotCallback(
- args.config,
- cfg,
- os.path.join(cfg.trial_dir, "configs"),
- use_version=False,
- ),
- ]
- if args.gradio:
- callbacks += [
- ProgressCallback(save_path=os.path.join(cfg.trial_dir, "progress"))
- ]
- else:
- callbacks += [CustomProgressBar(refresh_rate=1)]
-
- def write_to_text(file, lines):
- with open(file, "w") as f:
- for line in lines:
- f.write(line + "\n")
-
- loggers = []
- if args.train:
- # make tensorboard logging dir to suppress warning
- rank_zero_only(
- lambda: os.makedirs(os.path.join(cfg.trial_dir, "tb_logs"), exist_ok=True)
- )()
- loggers += [
- TensorBoardLogger(cfg.trial_dir, name="tb_logs"),
- CSVLogger(cfg.trial_dir, name="csv_logs"),
- ] + system.get_loggers()
- rank_zero_only(
- lambda: write_to_text(
- os.path.join(cfg.trial_dir, "cmd.txt"),
- ["python " + " ".join(sys.argv), str(args)],
- )
- )()
-
- trainer = Trainer(
- callbacks=callbacks,
- logger=loggers,
- inference_mode=False,
- accelerator="gpu",
- devices=devices,
- # profiler="advanced",
- **cfg.trainer,
- )
-
- def set_system_status(system: BaseSystem, ckpt_path: Optional[str]):
- if ckpt_path is None:
- return
- ckpt = torch.load(ckpt_path, map_location="cpu")
- system.set_resume_status(ckpt["epoch"], ckpt["global_step"])
- if args.train:
- trainer.fit(system, datamodule=dm, ckpt_path=cfg.resume)
- trainer.test(system, datamodule=dm)
- if args.gradio:
- # also export assets if in gradio mode
- trainer.predict(system, datamodule=dm)
- elif args.validate:
- # manually set epoch and global_step as they cannot be automatically resumed
- set_system_status(system, cfg.resume)
- trainer.validate(system, datamodule=dm, ckpt_path=cfg.resume)
- elif args.test:
- # manually set epoch and global_step as they cannot be automatically resumed
- set_system_status(system, cfg.resume)
- trainer.test(system, datamodule=dm, ckpt_path=cfg.resume)
- elif args.export:
- set_system_status(system, cfg.resume)
- trainer.predict(system, datamodule=dm, ckpt_path=cfg.resume)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument("--config", required=True, help="path to config file")
- parser.add_argument(
- "--gpu",
- default="0",
- help="GPU(s) to be used. 0 means use the 1st available GPU. "
- "1,2 means use the 2nd and 3rd available GPU. "
- "If CUDA_VISIBLE_DEVICES is set before calling `launch.py`, "
- "this argument is ignored and all available GPUs are always used.",
- )
-
- group = parser.add_mutually_exclusive_group(required=True)
- group.add_argument("--train", action="store_true")
- group.add_argument("--validate", action="store_true")
- group.add_argument("--test", action="store_true")
- group.add_argument("--export", action="store_true")
-
- parser.add_argument(
- "--gradio", action="store_true", help="if true, run in gradio mode"
- )
-
- parser.add_argument(
- "--verbose", action="store_true", help="if true, set logging level to DEBUG"
- )
-
- parser.add_argument(
- "--typecheck",
- action="store_true",
- help="whether to enable dynamic type checking",
- )
-
- args, extras = parser.parse_known_args()
-
- if args.gradio:
- with contextlib.redirect_stdout(sys.stderr):
- main(args, extras)
- else:
- main(args, extras)
diff --git a/packages.txt b/packages.txt
old mode 100644
new mode 100755
diff --git a/requirements.txt b/requirements.txt
old mode 100644
new mode 100755
-By mimicking the artist/craftsman modeling workflow, we propose CraftsMan (aka 匠心) which uses 3D Latent Set Diffusion Model that directly generates coarse meshes, +By mimicking the artist/craftsman modeling workflow, we propose CraftsMan (aka 匠心) that uses 3D Latent Set Diffusion Model that directly generate coarse meshes, then a multi-view normal enhanced image generation model is used to refine the mesh. We provide the coarse 3D diffusion part here.
@@ -36,8 +42,6 @@ If you found CraftsMan is helpful, please help to ⭐ the
*If you have your own multi-view images, you can directly upload it. -Tutorial -使用教程