test

.gitattributes

@@ -1,35 +0,0 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,155 +0,0 @@
-# Byte-compiled / optimized / DLL files
-__pycache__/
-*.py[cod]
-*$py.class
-
-# C extensions
-*.so
-
-# Distribution / packaging
-.Python
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-share/python-wheels/
-*.egg-info/
-.installed.cfg
-*.egg
-MANIFEST
-
-# PyInstaller
-#  Usually these files are written by a python script from a template
-#  before PyInstaller builds the exe, so as to inject date/other infos into it.
-*.manifest
-*.spec
-
-# Installer logs
-pip-log.txt
-pip-delete-this-directory.txt
-
-# Unit test / coverage reports
-htmlcov/
-.tox/
-.nox/
-.coverage
-.coverage.*
-.cache
-nosetests.xml
-coverage.xml
-*.cover
-*.py,cover
-.hypothesis/
-.pytest_cache/
-cover/
-
-# Translations
-*.mo
-*.pot
-
-# Django stuff:
-*.log
-local_settings.py
-db.sqlite3
-db.sqlite3-journal
-
-# Flask stuff:
-instance/
-.webassets-cache
-
-# Scrapy stuff:
-.scrapy
-
-# Sphinx documentation
-docs/_build/
-
-# PyBuilder
-.pybuilder/
-target/
-
-# Jupyter Notebook
-.ipynb_checkpoints
-
-# IPython
-profile_default/
-ipython_config.py
-
-# pyenv
-#   For a library or package, you might want to ignore these files since the code is
-#   intended to run in multiple environments; otherwise, check them in:
-# .python-version
-
-# pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-#   However, in case of collaboration, if having platform-specific dependencies or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don't work, or not
-#   install all needed dependencies.
-#Pipfile.lock
-
-# poetry
-#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
-#poetry.lock
-
-# pdm
-#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
-#pdm.lock
-#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
-#   in version control.
-#   https://pdm.fming.dev/#use-with-ide
-.pdm.toml
-
-# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
-__pypackages__/
-
-# Celery stuff
-celerybeat-schedule
-celerybeat.pid
-
-# SageMath parsed files
-*.sage.py
-
-# Environments
-.env
-.venv
-env/
-venv/
-ENV/
-env.bak/
-venv.bak/
-
-# Spyder project settings
-.spyderproject
-.spyproject
-
-# Rope project settings
-.ropeproject
-
-# mkdocs documentation
-/site
-
-# mypy
-.mypy_cache/
-.dmypy.json
-dmypy.json
-
-# Pyre type checker
-.pyre/
-
-# pytype static type analyzer
-.pytype/
-
-# Cython debug symbols
-cython_debug/
-
-out/

app.py

@@ -121,12 +121,14 @@ parser.add_argument(
     help="config for stage2",
 )
 
-parser.add_argument("--device", type=str, default="cuda")
+parser.add_argument("--device", type=str, default="cpu")
 args = parser.parse_args()
 
 crm_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="CRM.pth")
 specs = json.load(open("configs/specs_objaverse_total.json"))
-model = CRM(specs).to(args.device)
+# model = CRM(specs).to(args.device)
+model = CRM(specs).to("cpu")
+
 model.load_state_dict(torch.load(crm_path, map_location = args.device), strict=False)
 
 stage1_config = OmegaConf.load(args.stage1_config).config

imagedream/ldm/models/diffusion/ddim.py

@@ -22,8 +22,8 @@ class DDIMSampler(object):
 
     def register_buffer(self, name, attr):
         if type(attr) == torch.Tensor:
-            if attr.device != torch.device("cuda"):
-                attr = attr.to(torch.device("cuda"))
+            if attr.device != torch.device("cpu"):
+                attr = attr.to(torch.device("cpu"))
         setattr(self, name, attr)
 
     def make_schedule(

imagedream/ldm/modules/attention.py

@@ -226,6 +226,43 @@ class MemoryEfficientCrossAttention(nn.Module):
 
 
 class BasicTransformerBlock(nn.Module):
+    # def __init__(
+    #     self,
+    #     dim,
+    #     n_heads,
+    #     d_head,
+    #     dropout=0.0,
+    #     context_dim=None,
+    #     gated_ff=True,
+    #     checkpoint=True,
+    #     disable_self_attn=False,
+    #     **kwargs
+    # ):
+    #     super().__init__()
+    #     assert XFORMERS_IS_AVAILBLE, "xformers is not available"
+    #     attn_cls = MemoryEfficientCrossAttention
+    #     self.disable_self_attn = disable_self_attn
+    #     self.attn1 = attn_cls(
+    #         query_dim=dim,
+    #         heads=n_heads,
+    #         dim_head=d_head,
+    #         dropout=dropout,
+    #         context_dim=context_dim if self.disable_self_attn else None,
+    #     )  # is a self-attention if not self.disable_self_attn
+    #     self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+    #     self.attn2 = attn_cls(
+    #         query_dim=dim,
+    #         context_dim=context_dim,
+    #         heads=n_heads,
+    #         dim_head=d_head,
+    #         dropout=dropout,
+    #         **kwargs
+    #     )  # is self-attn if context is none
+    #     self.norm1 = nn.LayerNorm(dim)
+    #     self.norm2 = nn.LayerNorm(dim)
+    #     self.norm3 = nn.LayerNorm(dim)
+    #     self.checkpoint = checkpoint
+
     def __init__(
         self,
         dim,
@@ -239,7 +276,6 @@ class BasicTransformerBlock(nn.Module):
         **kwargs
     ):
         super().__init__()
-        assert XFORMERS_IS_AVAILBLE, "xformers is not available"
         attn_cls = MemoryEfficientCrossAttention
         self.disable_self_attn = disable_self_attn
         self.attn1 = attn_cls(
@@ -248,7 +284,7 @@ class BasicTransformerBlock(nn.Module):
             dim_head=d_head,
             dropout=dropout,
             context_dim=context_dim if self.disable_self_attn else None,
-        )  # is a self-attention if not self.disable_self_attn
+        )  # Self-attention if not self.disable_self_attn
         self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
         self.attn2 = attn_cls(
             query_dim=dim,
@@ -257,12 +293,13 @@ class BasicTransformerBlock(nn.Module):
             dim_head=d_head,
             dropout=dropout,
             **kwargs
-        )  # is self-attn if context is none
+        )  # Cross-attention if context is provided
         self.norm1 = nn.LayerNorm(dim)
         self.norm2 = nn.LayerNorm(dim)
         self.norm3 = nn.LayerNorm(dim)
         self.checkpoint = checkpoint
 
+
     def forward(self, x, context=None):
         return checkpoint(
             self._forward, (x, context), self.parameters(), self.checkpoint
@@ -278,6 +315,7 @@ class BasicTransformerBlock(nn.Module):
         x = self.attn2(self.norm2(x), context=context) + x
         x = self.ff(self.norm3(x)) + x
         return x
+    
 
 
 class SpatialTransformer(nn.Module):

imagedream/ldm/modules/diffusionmodules/model.py

@@ -220,39 +220,59 @@ class MemoryEfficientAttnBlock(nn.Module):
         self.attention_op: Optional[Any] = None
 
     def forward(self, x):
+        # h_ = x
+        # h_ = self.norm(h_)
+        # q = self.q(h_)
+        # k = self.k(h_)
+        # v = self.v(h_)
+
+        # # compute attention
+        # B, C, H, W = q.shape
+        # q, k, v = map(lambda x: rearrange(x, "b c h w -> b (h w) c"), (q, k, v))
+
+        # q, k, v = map(
+        #     lambda t: t.unsqueeze(3)
+        #     .reshape(B, t.shape[1], 1, C)
+        #     .permute(0, 2, 1, 3)
+        #     .reshape(B * 1, t.shape[1], C)
+        #     .contiguous(),
+        #     (q, k, v),
+        # )
+        # out = xformers.ops.memory_efficient_attention(
+        #     q, k, v, attn_bias=None, op=self.attention_op
+        # )
+
+        # out = (
+        #     out.unsqueeze(0)
+        #     .reshape(B, 1, out.shape[1], C)
+        #     .permute(0, 2, 1, 3)
+        #     .reshape(B, out.shape[1], C)
+        # )
+        # out = rearrange(out, "b (h w) c -> b c h w", b=B, h=H, w=W, c=C)
+        # out = self.proj_out(out)
+        # return x + out
         h_ = x
         h_ = self.norm(h_)
         q = self.q(h_)
         k = self.k(h_)
         v = self.v(h_)
 
-        # compute attention
+        # Compute attention
         B, C, H, W = q.shape
-        q, k, v = map(lambda x: rearrange(x, "b c h w -> b (h w) c"), (q, k, v))
+        q, k, v = map(lambda t: rearrange(t, "b c h w -> b (h w) c"), (q, k, v))
 
-        q, k, v = map(
-            lambda t: t.unsqueeze(3)
-            .reshape(B, t.shape[1], 1, C)
-            .permute(0, 2, 1, 3)
-            .reshape(B * 1, t.shape[1], C)
-            .contiguous(),
-            (q, k, v),
-        )
-        out = xformers.ops.memory_efficient_attention(
-            q, k, v, attn_bias=None, op=self.attention_op
-        )
+        if torch.cuda.is_available():  # Use xformers only if GPU is available
+            out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
+        else:
+            # CPU-friendly alternative for attention
+            attn_weights = torch.einsum('bqc,bkc->bqk', q, k)  # Simple dot-product attention
+            attn_weights = torch.softmax(attn_weights, dim=-1)
+            out = torch.einsum('bqk,bvc->bqc', attn_weights, v)
 
-        out = (
-            out.unsqueeze(0)
-            .reshape(B, 1, out.shape[1], C)
-            .permute(0, 2, 1, 3)
-            .reshape(B, out.shape[1], C)
-        )
-        out = rearrange(out, "b (h w) c -> b c h w", b=B, h=H, w=W, c=C)
+        out = rearrange(out, "b (h w) c -> b c h w", h=H, w=W)
         out = self.proj_out(out)
         return x + out
 
-
 class MemoryEfficientCrossAttentionWrapper(MemoryEfficientCrossAttention):
     def forward(self, x, context=None, mask=None):
         b, c, h, w = x.shape
@@ -263,6 +283,29 @@ class MemoryEfficientCrossAttentionWrapper(MemoryEfficientCrossAttention):
 
 
 def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None):
+    # assert attn_type in [
+    #     "vanilla",
+    #     "vanilla-xformers",
+    #     "memory-efficient-cross-attn",
+    #     "linear",
+    #     "none",
+    # ], f"attn_type {attn_type} unknown"
+    # if XFORMERS_IS_AVAILBLE and attn_type == "vanilla":
+    #     attn_type = "vanilla-xformers"
+    # print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
+    # if attn_type == "vanilla":
+    #     assert attn_kwargs is None
+    #     return AttnBlock(in_channels)
+    # elif attn_type == "vanilla-xformers":
+    #     print(f"building MemoryEfficientAttnBlock with {in_channels} in_channels...")
+    #     return MemoryEfficientAttnBlock(in_channels)
+    # elif type == "memory-efficient-cross-attn":
+    #     attn_kwargs["query_dim"] = in_channels
+    #     return MemoryEfficientCrossAttentionWrapper(**attn_kwargs)
+    # elif attn_type == "none":
+    #     return nn.Identity(in_channels)
+    # else:
+    #     raise NotImplementedError()
     assert attn_type in [
         "vanilla",
         "vanilla-xformers",
@@ -270,16 +313,22 @@ def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None):
         "linear",
         "none",
     ], f"attn_type {attn_type} unknown"
-    if XFORMERS_IS_AVAILBLE and attn_type == "vanilla":
+    
+    # Comprobar si GPU está disponible y evitar xformers si no lo está
+    if torch.cuda.is_available() and XFORMERS_IS_AVAILBLE and attn_type == "vanilla":
         attn_type = "vanilla-xformers"
+    else:
+        print("Using CPU-based attention as xformers or GPU is not available.")
+
     print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
+    
     if attn_type == "vanilla":
         assert attn_kwargs is None
-        return AttnBlock(in_channels)
+        return AttnBlock(in_channels)  # Atención estándar para CPU
     elif attn_type == "vanilla-xformers":
         print(f"building MemoryEfficientAttnBlock with {in_channels} in_channels...")
-        return MemoryEfficientAttnBlock(in_channels)
-    elif type == "memory-efficient-cross-attn":
+        return MemoryEfficientAttnBlock(in_channels)  # Atención optimizada con xformers
+    elif attn_type == "memory-efficient-cross-attn":
         attn_kwargs["query_dim"] = in_channels
         return MemoryEfficientCrossAttentionWrapper(**attn_kwargs)
     elif attn_type == "none":

imagedream/ldm/modules/encoders/modules.py

@@ -106,7 +106,7 @@ class FrozenCLIPEmbedder(AbstractEncoder):
     def __init__(
         self,
         version="openai/clip-vit-large-patch14",
-        device="cuda",
+        device="cpu",
         max_length=77,
         freeze=True,
         layer="last",

mesh.py

@@ -159,7 +159,7 @@ class Mesh:
 
         # device
         if device is None:
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            device = torch.device("cpu" if torch.cuda.is_available() else "cpu")
 
         mesh.device = device
 
@@ -331,7 +331,7 @@ class Mesh:
 
         # device
         if device is None:
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            device = torch.device("cpu" if torch.cuda.is_available() else "cpu")
 
         mesh.device = device
 

model/archs/unet.py

@@ -40,7 +40,7 @@ class UNetPP(nn.Module):
                 ),
             )
              
-        self.unet.enable_xformers_memory_efficient_attention()    
+        # self.unet.enable_xformers_memory_efficient_attention()    
         if in_channels > 12:
             self.learned_plane = torch.nn.parameter.Parameter(torch.zeros([1,in_channels-12,256,256*3]))
 

pipelines.py

@@ -15,7 +15,7 @@ class TwoStagePipeline(object):
         stage2_model_config,
         stage1_sampler_config,
         stage2_sampler_config,
-        device="cuda",
+        device="cpu",
         dtype=torch.float16,
         resize_rate=1,
     ) -> None:

run.py

@@ -125,8 +125,10 @@ if __name__ == "__main__":
 
     crm_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="CRM.pth")
     specs = json.load(open("configs/specs_objaverse_total.json"))
-    model = CRM(specs).to("cuda")
-    model.load_state_dict(torch.load(crm_path, map_location = "cuda"), strict=False)
+    # model = CRM(specs).to("cuda")
+    model = CRM(specs).to("cpu")
+
+    model.load_state_dict(torch.load(crm_path, map_location = "cpu"), strict=False)
 
     stage1_config = OmegaConf.load("configs/nf7_v3_SNR_rd_size_stroke.yaml").config
     stage2_config = OmegaConf.load("configs/stage2-v2-snr.yaml").config
@@ -156,5 +158,5 @@ if __name__ == "__main__":
     Image.fromarray(np_imgs).save(args.outdir+"pixel_images.png")
     Image.fromarray(np_xyzs).save(args.outdir+"xyz_images.png")
 
-    glb_path, obj_path = generate3d(model, np_imgs, np_xyzs, "cuda")
+    glb_path, obj_path = generate3d(model, np_imgs, np_xyzs, "cpu")
     shutil.copy(obj_path, args.outdir+"output3d.zip")

util/flexicubes.py

@@ -64,7 +64,7 @@ class FlexiCubes:
             The scale of weights in FlexiCubes. Should be between 0 and 1.
     """
 
-    def __init__(self, device="cuda", qef_reg_scale=1e-3, weight_scale=0.99):
+    def __init__(self, device="cpu", qef_reg_scale=1e-3, weight_scale=0.99):
 
         self.device = device
         self.dmc_table = torch.tensor(dmc_table, dtype=torch.long, device=device, requires_grad=False)

util/flexicubes_geometry.py

@@ -31,7 +31,7 @@ def get_center_boundary_index(grid_res, device):
 ###############################################################################
 class FlexiCubesGeometry(object):
     def __init__(
-            self, grid_res=64, scale=2.0, device='cuda', renderer=None,
+            self, grid_res=64, scale=2.0, device='cpu', renderer=None,
             render_type='neural_render', args=None):
         super(FlexiCubesGeometry, self).__init__()
         self.grid_res = grid_res

util/renderer.py

@@ -1,7 +1,56 @@
 
+# import torch
+# import torch.nn as nn
+# import nvdiffrast.torch as dr
+# from util.flexicubes_geometry import FlexiCubesGeometry
+
+# class Renderer(nn.Module):
+#     def __init__(self, tet_grid_size, camera_angle_num, scale, geo_type):
+#         super().__init__()
+
+#         self.tet_grid_size = tet_grid_size
+#         self.camera_angle_num = camera_angle_num
+#         self.scale = scale
+#         self.geo_type = geo_type
+#         self.glctx = dr.RasterizeCudaContext()
+
+#         if self.geo_type == "flex":
+#             self.flexicubes = FlexiCubesGeometry(grid_res = self.tet_grid_size)   
+
+#     def forward(self, data, sdf, deform, verts, tets, training=False, weight = None):
+
+#         results = {}
+
+#         deform = torch.tanh(deform) / self.tet_grid_size * self.scale / 0.95
+#         if self.geo_type == "flex":
+#             deform = deform *0.5
+
+#             v_deformed = verts + deform
+
+#             verts_list = []
+#             faces_list = []
+#             reg_list = []
+#             n_shape = verts.shape[0]
+#             for i in range(n_shape): 
+#                 verts_i, faces_i, reg_i = self.flexicubes.get_mesh(v_deformed[i], sdf[i].squeeze(dim=-1),
+#                 with_uv=False, indices=tets, weight_n=weight[i], is_training=training)
+
+#                 verts_list.append(verts_i)
+#                 faces_list.append(faces_i)
+#                 reg_list.append(reg_i)       
+#             verts = verts_list
+#             faces = faces_list
+
+#             flexicubes_surface_reg = torch.cat(reg_list).mean()
+#             flexicubes_weight_reg = (weight ** 2).mean()
+#             results["flex_surf_loss"] = flexicubes_surface_reg
+#             results["flex_weight_loss"] = flexicubes_weight_reg
+
+#         return results, verts, faces
+
 import torch
 import torch.nn as nn
-import nvdiffrast.torch as dr
+# import nvdiffrast.torch as dr  # Comentado porque no se usará en CPU
 from util.flexicubes_geometry import FlexiCubesGeometry
 
 class Renderer(nn.Module):
@@ -12,18 +61,19 @@ class Renderer(nn.Module):
         self.camera_angle_num = camera_angle_num
         self.scale = scale
         self.geo_type = geo_type
-        self.glctx = dr.RasterizeCudaContext()
+        
+        # Eliminar el contexto de GPU y usar una alternativa o desactivarlo
+        # self.glctx = dr.RasterizeCudaContext()  # Comentado porque se usa GPU
 
         if self.geo_type == "flex":
-            self.flexicubes = FlexiCubesGeometry(grid_res = self.tet_grid_size)   
-
-    def forward(self, data, sdf, deform, verts, tets, training=False, weight = None):
+            self.flexicubes = FlexiCubesGeometry(grid_res=self.tet_grid_size)
 
+    def forward(self, data, sdf, deform, verts, tets, training=False, weight=None):
         results = {}
 
         deform = torch.tanh(deform) / self.tet_grid_size * self.scale / 0.95
         if self.geo_type == "flex":
-            deform = deform *0.5
+            deform = deform * 0.5
 
             v_deformed = verts + deform
 
@@ -31,13 +81,17 @@ class Renderer(nn.Module):
             faces_list = []
             reg_list = []
             n_shape = verts.shape[0]
-            for i in range(n_shape): 
-                verts_i, faces_i, reg_i = self.flexicubes.get_mesh(v_deformed[i], sdf[i].squeeze(dim=-1),
-                with_uv=False, indices=tets, weight_n=weight[i], is_training=training)
+            for i in range(n_shape):
+                # Aquí deberás adaptar el uso de FlexiCubesGeometry para que funcione sin GPU.
+                verts_i, faces_i, reg_i = self.flexicubes.get_mesh(
+                    v_deformed[i], sdf[i].squeeze(dim=-1),
+                    with_uv=False, indices=tets, weight_n=weight[i], is_training=training
+                )
 
                 verts_list.append(verts_i)
                 faces_list.append(faces_i)
-                reg_list.append(reg_i)       
+                reg_list.append(reg_i)
+                
             verts = verts_list
             faces = faces_list
 
@@ -46,4 +100,4 @@ class Renderer(nn.Module):
             results["flex_surf_loss"] = flexicubes_surface_reg
             results["flex_weight_loss"] = flexicubes_weight_reg
 
-        return results, verts, faces
+        return results, verts, faces