BladeSzaSza commited on
Commit
2153bff
Β·
1 Parent(s): 0b42bb7

added logs

Browse files
Files changed (3) hide show
  1. app.py +3 -0
  2. models/image_generator.py +5 -6
  3. models/model_3d_generator.py +144 -23
app.py CHANGED
@@ -7,6 +7,9 @@ import gc
7
  from datetime import datetime
8
  from pathlib import Path
9
 
 
 
 
10
  # Initialize directories
11
  DATA_DIR = Path("/data") if os.path.exists("/data") else Path("./data")
12
  DATA_DIR.mkdir(exist_ok=True)
 
7
  from datetime import datetime
8
  from pathlib import Path
9
 
10
+ # Disable torch dynamo globally to avoid ConstantVariable errors
11
+ torch._dynamo.config.suppress_errors = True
12
+
13
  # Initialize directories
14
  DATA_DIR = Path("/data") if os.path.exists("/data") else Path("./data")
15
  DATA_DIR.mkdir(exist_ok=True)
models/image_generator.py CHANGED
@@ -5,6 +5,9 @@ import numpy as np
5
  from typing import Optional, List, Union
6
  import gc
7
 
 
 
 
8
  class OmniGenImageGenerator:
9
  """Image generation using OmniGen2 model"""
10
 
@@ -64,12 +67,8 @@ class OmniGenImageGenerator:
64
  else:
65
  self.pipeline = self.pipeline.to(self.device)
66
 
67
- # Compile for faster inference (if available)
68
- if hasattr(torch, 'compile') and self.device == "cuda":
69
- try:
70
- self.pipeline.unet = torch.compile(self.pipeline.unet, mode="reduce-overhead")
71
- except:
72
- pass # Compilation is optional
73
 
74
  except Exception as e:
75
  print(f"Failed to load image generation model: {e}")
 
5
  from typing import Optional, List, Union
6
  import gc
7
 
8
+ # Disable torch dynamo to avoid ConstantVariable errors
9
+ torch._dynamo.config.suppress_errors = True
10
+
11
  class OmniGenImageGenerator:
12
  """Image generation using OmniGen2 model"""
13
 
 
67
  else:
68
  self.pipeline = self.pipeline.to(self.device)
69
 
70
+ # Disable torch.compile to avoid dynamo issues that cause ConstantVariable errors
71
+ print("Skipping torch.compile to avoid dynamo compatibility issues")
 
 
 
 
72
 
73
  except Exception as e:
74
  print(f"Failed to load image generation model: {e}")
models/model_3d_generator.py CHANGED
@@ -6,12 +6,21 @@ import tempfile
6
  from typing import Union, Optional, Dict, Any
7
  from pathlib import Path
8
  import os
 
 
 
 
 
9
 
10
  class Hunyuan3DGenerator:
11
  """3D model generation using Hunyuan3D-2.1"""
12
 
13
  def __init__(self, device: str = "cuda"):
 
 
14
  self.device = device if torch.cuda.is_available() else "cpu"
 
 
15
  self.model = None
16
  self.preprocessor = None
17
 
@@ -25,52 +34,106 @@ class Hunyuan3DGenerator:
25
  self.resolution = 256 # 3D resolution
26
 
27
  # Use lite model for low VRAM
28
- self.use_lite = self.device == "cpu" or not self._check_vram()
 
 
 
 
29
 
30
  def _check_vram(self) -> bool:
31
  """Check if we have enough VRAM for full model"""
 
 
32
  if not torch.cuda.is_available():
 
33
  return False
34
 
35
  try:
36
  vram = torch.cuda.get_device_properties(0).total_memory
 
 
 
37
  # Need at least 12GB for full model
38
- return vram > 12 * 1024 * 1024 * 1024
39
- except:
 
 
 
40
  return False
41
 
42
  def load_model(self):
43
  """Lazy load the 3D generation model"""
44
  if self.model is None:
 
 
45
  try:
46
  # Import Hunyuan3D components
 
47
  from transformers import AutoModel, AutoProcessor
48
 
49
  model_id = self.lite_model_id if self.use_lite else self.model_id
 
50
 
51
  # Load preprocessor
 
52
  self.preprocessor = AutoProcessor.from_pretrained(model_id)
 
53
 
54
  # Load model with optimizations
55
  torch_dtype = torch.float16 if self.device == "cuda" else torch.float32
 
56
 
 
 
 
 
 
57
  self.model = AutoModel.from_pretrained(
58
  model_id,
59
  torch_dtype=torch_dtype,
60
  low_cpu_mem_usage=True,
61
- device_map="auto" if self.device == "cuda" else None,
62
  trust_remote_code=True
63
  )
 
64
 
65
- if self.device == "cpu":
66
- self.model = self.model.to(self.device)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
67
 
68
- # Enable optimizations
 
69
  if hasattr(self.model, 'enable_attention_slicing'):
70
  self.model.enable_attention_slicing()
 
 
 
 
 
71
 
72
  except Exception as e:
73
- print(f"Failed to load Hunyuan3D model: {e}")
 
 
74
  # Model loading failed, will use fallback
75
  self.model = "fallback"
76
 
@@ -80,55 +143,113 @@ class Hunyuan3DGenerator:
80
  texture_resolution: int = 1024) -> Union[str, trimesh.Trimesh]:
81
  """Convert 2D image to 3D model"""
82
 
 
 
 
 
 
83
  try:
84
  # Load model if needed
 
85
  if self.model is None:
 
86
  self.load_model()
 
 
87
 
88
  # If model loading failed, use fallback
89
  if self.model == "fallback":
 
90
  return self._generate_fallback_3d(image)
91
 
92
  # Prepare image
 
93
  if isinstance(image, str):
 
94
  image = Image.open(image)
95
  elif isinstance(image, np.ndarray):
 
96
  image = Image.fromarray(image)
 
 
97
 
98
  # Ensure RGB
 
99
  if image.mode != 'RGB':
 
100
  image = image.convert('RGB')
101
 
 
 
102
  # Resize for processing
 
103
  image = image.resize((512, 512), Image.Resampling.LANCZOS)
 
104
 
105
  # Remove background if requested
106
  if remove_background:
107
- image = self._remove_background(image)
 
 
 
 
 
 
108
 
109
  # Process with model
 
110
  with torch.no_grad():
111
- # Preprocess image
112
- inputs = self.preprocessor(images=image, return_tensors="pt").to(self.device)
113
-
114
- # Generate 3D
115
- outputs = self.model.generate(
116
- **inputs,
117
- num_inference_steps=self.num_inference_steps,
118
- guidance_scale=self.guidance_scale,
119
- texture_resolution=texture_resolution
120
- )
121
-
122
- # Extract mesh
123
- mesh = self._extract_mesh(outputs)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
124
 
125
  # Save mesh
 
126
  mesh_path = self._save_mesh(mesh)
 
127
 
 
128
  return mesh_path
129
 
130
  except Exception as e:
131
- print(f"3D generation error: {e}")
 
 
132
  return self._generate_fallback_3d(image)
133
 
134
  def _remove_background(self, image: Image.Image) -> Image.Image:
 
6
  from typing import Union, Optional, Dict, Any
7
  from pathlib import Path
8
  import os
9
+ import logging
10
+
11
+ # Set up detailed logging for 3D generation
12
+ logging.basicConfig(level=logging.INFO)
13
+ logger = logging.getLogger(__name__)
14
 
15
  class Hunyuan3DGenerator:
16
  """3D model generation using Hunyuan3D-2.1"""
17
 
18
  def __init__(self, device: str = "cuda"):
19
+ logger.info(f"πŸ”§ Initializing Hunyuan3DGenerator with device: {device}")
20
+
21
  self.device = device if torch.cuda.is_available() else "cpu"
22
+ logger.info(f"πŸ”§ Final device selection: {self.device}")
23
+
24
  self.model = None
25
  self.preprocessor = None
26
 
 
34
  self.resolution = 256 # 3D resolution
35
 
36
  # Use lite model for low VRAM
37
+ vram_check = self._check_vram()
38
+ self.use_lite = self.device == "cpu" or not vram_check
39
+
40
+ logger.info(f"πŸ”§ VRAM check result: {vram_check}, using lite model: {self.use_lite}")
41
+ logger.info(f"πŸ”§ Model ID to use: {self.lite_model_id if self.use_lite else self.model_id}")
42
 
43
  def _check_vram(self) -> bool:
44
  """Check if we have enough VRAM for full model"""
45
+ logger.info("πŸ” Checking VRAM availability...")
46
+
47
  if not torch.cuda.is_available():
48
+ logger.info("❌ CUDA not available")
49
  return False
50
 
51
  try:
52
  vram = torch.cuda.get_device_properties(0).total_memory
53
+ vram_gb = vram / (1024 * 1024 * 1024)
54
+ logger.info(f"πŸ” Available VRAM: {vram_gb:.2f} GB")
55
+
56
  # Need at least 12GB for full model
57
+ has_enough = vram > 12 * 1024 * 1024 * 1024
58
+ logger.info(f"πŸ” Has enough VRAM (>12GB): {has_enough}")
59
+ return has_enough
60
+ except Exception as e:
61
+ logger.error(f"❌ Error checking VRAM: {e}")
62
  return False
63
 
64
  def load_model(self):
65
  """Lazy load the 3D generation model"""
66
  if self.model is None:
67
+ logger.info("πŸš€ Starting 3D model loading process...")
68
+
69
  try:
70
  # Import Hunyuan3D components
71
+ logger.info("πŸ“¦ Importing transformers components...")
72
  from transformers import AutoModel, AutoProcessor
73
 
74
  model_id = self.lite_model_id if self.use_lite else self.model_id
75
+ logger.info(f"πŸ“¦ Loading model: {model_id}")
76
 
77
  # Load preprocessor
78
+ logger.info("πŸ“¦ Loading preprocessor...")
79
  self.preprocessor = AutoProcessor.from_pretrained(model_id)
80
+ logger.info("βœ… Preprocessor loaded successfully")
81
 
82
  # Load model with optimizations
83
  torch_dtype = torch.float16 if self.device == "cuda" else torch.float32
84
+ logger.info(f"πŸ“¦ Using torch dtype: {torch_dtype}")
85
 
86
+ # Disable torch.compile to avoid dynamo issues
87
+ logger.info("πŸ“¦ Disabling torch compile to avoid dynamo issues...")
88
+ torch._dynamo.config.suppress_errors = True
89
+
90
+ logger.info("πŸ“¦ Loading 3D model with safe device handling...")
91
  self.model = AutoModel.from_pretrained(
92
  model_id,
93
  torch_dtype=torch_dtype,
94
  low_cpu_mem_usage=True,
95
+ device_map=None, # Avoid auto device mapping to prevent meta tensor issues
96
  trust_remote_code=True
97
  )
98
+ logger.info("βœ… 3D model loaded from pretrained")
99
 
100
+ # Safe device movement
101
+ logger.info(f"πŸ“¦ Moving model to device: {self.device}")
102
+ try:
103
+ if hasattr(self.model, 'to_empty'):
104
+ # Use to_empty for meta tensors
105
+ logger.info("πŸ“¦ Using to_empty() for safe device movement...")
106
+ self.model = self.model.to_empty(device=self.device, dtype=torch_dtype)
107
+ else:
108
+ # Standard device movement
109
+ logger.info("πŸ“¦ Using standard to() for device movement...")
110
+ self.model = self.model.to(self.device, dtype=torch_dtype)
111
+ logger.info("βœ… Model successfully moved to device")
112
+
113
+ except RuntimeError as device_error:
114
+ logger.error(f"❌ Device movement failed: {device_error}")
115
+ if "meta tensor" in str(device_error):
116
+ logger.info("πŸ”„ Attempting CPU fallback for meta tensor issue...")
117
+ self.device = "cpu"
118
+ self.model = self.model.to("cpu")
119
+ logger.info("βœ… Fallback to CPU successful")
120
+ else:
121
+ raise device_error
122
 
123
+ # Enable optimizations safely
124
+ logger.info("πŸ“¦ Applying model optimizations...")
125
  if hasattr(self.model, 'enable_attention_slicing'):
126
  self.model.enable_attention_slicing()
127
+ logger.info("βœ… Attention slicing enabled")
128
+ else:
129
+ logger.info("⚠️ Attention slicing not available")
130
+
131
+ logger.info("πŸŽ‰ 3D model loading completed successfully!")
132
 
133
  except Exception as e:
134
+ logger.error(f"❌ Failed to load Hunyuan3D model: {e}")
135
+ logger.error(f"❌ Error type: {type(e).__name__}")
136
+ logger.info("πŸ”„ Falling back to simple 3D generation...")
137
  # Model loading failed, will use fallback
138
  self.model = "fallback"
139
 
 
143
  texture_resolution: int = 1024) -> Union[str, trimesh.Trimesh]:
144
  """Convert 2D image to 3D model"""
145
 
146
+ logger.info("🎯 Starting image-to-3D conversion process...")
147
+ logger.info(f"🎯 Input type: {type(image)}")
148
+ logger.info(f"🎯 Remove background: {remove_background}")
149
+ logger.info(f"🎯 Texture resolution: {texture_resolution}")
150
+
151
  try:
152
  # Load model if needed
153
+ logger.info("πŸ” Checking if model needs loading...")
154
  if self.model is None:
155
+ logger.info("πŸ“¦ Model not loaded, initiating loading...")
156
  self.load_model()
157
+ else:
158
+ logger.info("βœ… Model already loaded")
159
 
160
  # If model loading failed, use fallback
161
  if self.model == "fallback":
162
+ logger.info("πŸ”„ Using fallback 3D generation...")
163
  return self._generate_fallback_3d(image)
164
 
165
  # Prepare image
166
+ logger.info("πŸ–ΌοΈ Preparing input image...")
167
  if isinstance(image, str):
168
+ logger.info(f"πŸ–ΌοΈ Loading image from path: {image}")
169
  image = Image.open(image)
170
  elif isinstance(image, np.ndarray):
171
+ logger.info("πŸ–ΌοΈ Converting numpy array to PIL Image")
172
  image = Image.fromarray(image)
173
+ else:
174
+ logger.info("πŸ–ΌοΈ Input is already PIL Image")
175
 
176
  # Ensure RGB
177
+ logger.info(f"πŸ–ΌοΈ Image mode: {image.mode}")
178
  if image.mode != 'RGB':
179
+ logger.info("πŸ–ΌοΈ Converting image to RGB mode")
180
  image = image.convert('RGB')
181
 
182
+ logger.info(f"πŸ–ΌοΈ Final image size: {image.size}")
183
+
184
  # Resize for processing
185
+ logger.info("πŸ–ΌοΈ Resizing image for processing (512x512)...")
186
  image = image.resize((512, 512), Image.Resampling.LANCZOS)
187
+ logger.info("βœ… Image resized successfully")
188
 
189
  # Remove background if requested
190
  if remove_background:
191
+ logger.info("🎭 Removing background from image...")
192
+ try:
193
+ image = self._remove_background(image)
194
+ logger.info("βœ… Background removal completed")
195
+ except Exception as bg_error:
196
+ logger.error(f"❌ Background removal failed: {bg_error}")
197
+ logger.info("πŸ”„ Continuing with original image...")
198
 
199
  # Process with model
200
+ logger.info("🧠 Starting model inference...")
201
  with torch.no_grad():
202
+ try:
203
+ # Preprocess image
204
+ logger.info("πŸ”„ Preprocessing image for model...")
205
+ inputs = self.preprocessor(images=image, return_tensors="pt")
206
+ logger.info(f"πŸ”„ Input tensor shape: {inputs['pixel_values'].shape if 'pixel_values' in inputs else 'unknown'}")
207
+
208
+ # Move inputs to device safely
209
+ logger.info(f"πŸ”„ Moving inputs to device: {self.device}")
210
+ try:
211
+ # Avoid device-related dynamo issues
212
+ device_str = str(self.device) # Convert to string to avoid torch.device in dynamo
213
+ inputs = {k: v.to(device_str) for k, v in inputs.items() if hasattr(v, 'to')}
214
+ logger.info("βœ… Inputs moved to device successfully")
215
+ except Exception as device_error:
216
+ logger.error(f"❌ Failed to move inputs to device: {device_error}")
217
+ raise device_error
218
+
219
+ # Generate 3D
220
+ logger.info("πŸš€ Starting 3D generation inference...")
221
+ logger.info(f"πŸš€ Parameters: steps={self.num_inference_steps}, guidance={self.guidance_scale}")
222
+
223
+ outputs = self.model.generate(
224
+ **inputs,
225
+ num_inference_steps=self.num_inference_steps,
226
+ guidance_scale=self.guidance_scale,
227
+ texture_resolution=texture_resolution
228
+ )
229
+ logger.info("βœ… 3D generation completed successfully")
230
+
231
+ # Extract mesh
232
+ logger.info("πŸ”§ Extracting mesh from model outputs...")
233
+ mesh = self._extract_mesh(outputs)
234
+ logger.info("βœ… Mesh extraction completed")
235
+
236
+ except Exception as inference_error:
237
+ logger.error(f"❌ Model inference failed: {inference_error}")
238
+ logger.error(f"❌ Inference error type: {type(inference_error).__name__}")
239
+ raise inference_error
240
 
241
  # Save mesh
242
+ logger.info("πŸ’Ύ Saving generated mesh...")
243
  mesh_path = self._save_mesh(mesh)
244
+ logger.info(f"βœ… Mesh saved to: {mesh_path}")
245
 
246
+ logger.info("πŸŽ‰ 3D generation process completed successfully!")
247
  return mesh_path
248
 
249
  except Exception as e:
250
+ logger.error(f"❌ 3D generation error: {e}")
251
+ logger.error(f"❌ Error type: {type(e).__name__}")
252
+ logger.info("πŸ”„ Falling back to simple 3D generation...")
253
  return self._generate_fallback_3d(image)
254
 
255
  def _remove_background(self, image: Image.Image) -> Image.Image: