YiftachEde commited on
Commit
49f568d
·
1 Parent(s): fcc9ef6
Files changed (1) hide show
  1. app.py +10 -28
app.py CHANGED
@@ -231,25 +231,14 @@ class ShapERenderer:
231
  def __init__(self, device):
232
  print("Initializing Shap-E models...")
233
  self.device = device
234
- self.xm = None
235
- self.model = None
236
- self.diffusion = None
 
237
  print("Shap-E models initialized!")
238
 
239
- def ensure_models_loaded(self):
240
- if self.xm is None:
241
- try:
242
- torch.cuda.empty_cache() # Clear GPU memory before loading
243
- self.xm = load_model('transmitter', device=self.device)
244
- self.model = load_model('text300M', device=self.device)
245
- self.diffusion = diffusion_from_config(load_config('diffusion'))
246
- except Exception as e:
247
- print(f"Error loading models: {e}")
248
- raise
249
-
250
  def generate_views(self, prompt, guidance_scale=15.0, num_steps=64):
251
  try:
252
- self.ensure_models_loaded()
253
  torch.cuda.empty_cache() # Clear GPU memory before generation
254
 
255
  # Generate latents using the text-to-3D model
@@ -313,24 +302,13 @@ class ShapERenderer:
313
  class RefinerInterface:
314
  def __init__(self):
315
  print("Initializing InstantMesh models...")
316
- self.pipeline = None
317
- self.model = None
318
- self.infer_config = None
319
  print("InstantMesh models initialized!")
320
 
321
- def ensure_models_loaded(self):
322
- if self.pipeline is None:
323
- try:
324
- torch.cuda.empty_cache() # Clear GPU memory before loading
325
- self.pipeline, self.model, self.infer_config = load_models()
326
- except Exception as e:
327
- print(f"Error loading models: {e}")
328
- raise
329
-
330
  def refine_model(self, input_image, prompt, steps=75, guidance_scale=7.5):
331
  """Main refinement function"""
332
  try:
333
- self.ensure_models_loaded()
334
  torch.cuda.empty_cache() # Clear GPU memory before processing
335
 
336
  # Process image and get refined output
@@ -404,9 +382,13 @@ class RefinerInterface:
404
  raise
405
 
406
  def create_demo():
 
407
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 
408
  shap_e = ShapERenderer(device)
409
  refiner = RefinerInterface()
 
410
 
411
  with gr.Blocks() as demo:
412
  gr.Markdown("# Shap-E to InstantMesh Pipeline")
 
231
  def __init__(self, device):
232
  print("Initializing Shap-E models...")
233
  self.device = device
234
+ torch.cuda.empty_cache() # Clear GPU memory before loading
235
+ self.xm = load_model('transmitter', device=self.device)
236
+ self.model = load_model('text300M', device=self.device)
237
+ self.diffusion = diffusion_from_config(load_config('diffusion'))
238
  print("Shap-E models initialized!")
239
 
 
 
 
 
 
 
 
 
 
 
 
240
  def generate_views(self, prompt, guidance_scale=15.0, num_steps=64):
241
  try:
 
242
  torch.cuda.empty_cache() # Clear GPU memory before generation
243
 
244
  # Generate latents using the text-to-3D model
 
302
  class RefinerInterface:
303
  def __init__(self):
304
  print("Initializing InstantMesh models...")
305
+ torch.cuda.empty_cache() # Clear GPU memory before loading
306
+ self.pipeline, self.model, self.infer_config = load_models()
 
307
  print("InstantMesh models initialized!")
308
 
 
 
 
 
 
 
 
 
 
309
  def refine_model(self, input_image, prompt, steps=75, guidance_scale=7.5):
310
  """Main refinement function"""
311
  try:
 
312
  torch.cuda.empty_cache() # Clear GPU memory before processing
313
 
314
  # Process image and get refined output
 
382
  raise
383
 
384
  def create_demo():
385
+ print("Initializing models...")
386
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
387
+
388
+ # Initialize models at startup
389
  shap_e = ShapERenderer(device)
390
  refiner = RefinerInterface()
391
+ print("All models initialized!")
392
 
393
  with gr.Blocks() as demo:
394
  gr.Markdown("# Shap-E to InstantMesh Pipeline")