mike23415 commited on
Commit
6740944
·
verified ·
1 Parent(s): b8e80f2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +30 -6
app.py CHANGED
@@ -34,10 +34,21 @@ def initialize_enhancer():
34
  # Force CPU usage for Hugging Face compatibility
35
  device = torch.device('cpu')
36
 
 
 
 
 
 
 
 
 
 
 
 
37
  # Initialize the upsampler
38
  upsampler = RealESRGANer(
39
  scale=4,
40
- model_path='weights/realesr-general-x4v3.pth',
41
  model=model,
42
  tile=0, # Set to 0 for small images, increase for large images
43
  tile_pad=10,
@@ -58,11 +69,14 @@ def initialize_enhancer():
58
  print(error_msg)
59
  return False
60
 
 
 
 
61
  # Initialize model immediately at startup - this blocks until model is ready
62
  if not initialize_enhancer():
63
- print("ERROR: Model failed to initialize. Server will exit.")
64
- # In production, you might prefer to exit when model fails
65
- # os.exit(1)
66
 
67
  @app.route('/enhance', methods=['POST'])
68
  def enhance_image():
@@ -110,10 +124,20 @@ def enhance_image():
110
 
111
  @app.route('/health', methods=['GET'])
112
  def health_check():
113
- global upsampler
 
 
 
 
 
 
 
 
 
 
114
 
115
  status_info = {
116
- 'status': 'ready' if upsampler is not None else 'failed',
117
  'timestamp': time.time()
118
  }
119
 
 
34
  # Force CPU usage for Hugging Face compatibility
35
  device = torch.device('cpu')
36
 
37
+ # Check if model weights file exists
38
+ weights_path = 'weights/realesr-general-x4v3.pth'
39
+ if not os.path.exists(weights_path):
40
+ print(f"Model weights not found at {weights_path}")
41
+ # Create the directory if it doesn't exist
42
+ os.makedirs('weights', exist_ok=True)
43
+ # You'd normally download weights here, but for this example
44
+ # we'll just use a placeholder
45
+ print("ERROR: Model weights file not found!")
46
+ return False
47
+
48
  # Initialize the upsampler
49
  upsampler = RealESRGANer(
50
  scale=4,
51
+ model_path=weights_path,
52
  model=model,
53
  tile=0, # Set to 0 for small images, increase for large images
54
  tile_pad=10,
 
69
  print(error_msg)
70
  return False
71
 
72
+ # Global init flag to track if we've attempted initialization
73
+ init_attempted = False
74
+
75
  # Initialize model immediately at startup - this blocks until model is ready
76
  if not initialize_enhancer():
77
+ print("ERROR: Model failed to initialize. Server will continue running but enhancement won't work.")
78
+ # We'll keep running but mark that we attempted initialization
79
+ init_attempted = True
80
 
81
  @app.route('/enhance', methods=['POST'])
82
  def enhance_image():
 
124
 
125
  @app.route('/health', methods=['GET'])
126
  def health_check():
127
+ global upsampler, init_attempted
128
+
129
+ # If we have the upsampler, we're ready
130
+ if upsampler is not None:
131
+ status = 'ready'
132
+ # If we tried to initialize and failed, report failure
133
+ elif init_attempted:
134
+ status = 'failed'
135
+ # Otherwise we're still in an unknown state
136
+ else:
137
+ status = 'initializing'
138
 
139
  status_info = {
140
+ 'status': status,
141
  'timestamp': time.time()
142
  }
143