gokaygokay commited on
Commit
2694d70
1 Parent(s): 8a735a9

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -9
app.py CHANGED
@@ -107,9 +107,8 @@ class ModelManager:
107
  self.realesrgan_x2 = None
108
  self.realesrgan_x4 = None
109
 
110
- def load_models(self, progress=gr.Progress()):
111
- if self.pipe is None:
112
- progress(0, desc="Loading Stable Diffusion pipeline...")
113
  self.pipe = self.setup_pipeline()
114
  self.pipe.to(device)
115
  self.pipe.unet.set_attn_processor(AttnProcessor2_0())
@@ -118,17 +117,14 @@ class ModelManager:
118
  progress(0.5, desc="Compiling the model...")
119
  self.pipe.unet = torch.compile(self.pipe.unet, mode="reduce-overhead", fullgraph=True)
120
 
121
- if self.realesrgan_x2 is None:
122
- progress(0.7, desc="Loading RealESRGAN x2 model...")
123
  self.realesrgan_x2 = RealESRGAN(device, scale=2)
124
  self.realesrgan_x2.load_weights('models/upscalers/RealESRGAN_x2.pth', download=False)
125
 
126
- if self.realesrgan_x4 is None:
127
- progress(0.9, desc="Loading RealESRGAN x4 model...")
128
  self.realesrgan_x4 = RealESRGAN(device, scale=4)
129
  self.realesrgan_x4.load_weights('models/upscalers/RealESRGAN_x4.pth', download=False)
130
 
131
- progress(1.0, desc="All models loaded successfully")
132
 
133
  def setup_pipeline(self):
134
  controlnet = ControlNetModel.from_single_file(
@@ -233,6 +229,7 @@ class ModelManager:
233
  return hdr_result
234
 
235
  model_manager = ModelManager()
 
236
 
237
  def extract_frames(video_path, output_folder):
238
  os.makedirs(output_folder, exist_ok=True)
@@ -274,11 +271,12 @@ def frames_to_video(input_folder, output_path, fps, original_video_path):
274
  # Remove the temporary file
275
  os.remove(temp_output_path)
276
 
 
277
  @timer_func
278
  def process_video(input_video, resolution, num_inference_steps, strength, hdr, guidance_scale, max_frames=None, frame_interval=1, preserve_frames=False, batch_size=8, progress=gr.Progress()):
279
  abort_event.clear() # Clear the abort flag at the start of a new job
280
  print("Starting video processing...")
281
- model_manager.load_models(progress) # Ensure models are loaded
282
 
283
  # Create a new job folder
284
  job_id = str(uuid.uuid4())
 
107
  self.realesrgan_x2 = None
108
  self.realesrgan_x4 = None
109
 
110
+ def load_models(self):
111
+ if self.pipe is None:
 
112
  self.pipe = self.setup_pipeline()
113
  self.pipe.to(device)
114
  self.pipe.unet.set_attn_processor(AttnProcessor2_0())
 
117
  progress(0.5, desc="Compiling the model...")
118
  self.pipe.unet = torch.compile(self.pipe.unet, mode="reduce-overhead", fullgraph=True)
119
 
120
+ if self.realesrgan_x2 is None:
 
121
  self.realesrgan_x2 = RealESRGAN(device, scale=2)
122
  self.realesrgan_x2.load_weights('models/upscalers/RealESRGAN_x2.pth', download=False)
123
 
124
+ if self.realesrgan_x4 is None:
 
125
  self.realesrgan_x4 = RealESRGAN(device, scale=4)
126
  self.realesrgan_x4.load_weights('models/upscalers/RealESRGAN_x4.pth', download=False)
127
 
 
128
 
129
  def setup_pipeline(self):
130
  controlnet = ControlNetModel.from_single_file(
 
229
  return hdr_result
230
 
231
  model_manager = ModelManager()
232
+ model_manager.load_models() # Ensure models are loaded
233
 
234
  def extract_frames(video_path, output_folder):
235
  os.makedirs(output_folder, exist_ok=True)
 
271
  # Remove the temporary file
272
  os.remove(temp_output_path)
273
 
274
+
275
  @timer_func
276
  def process_video(input_video, resolution, num_inference_steps, strength, hdr, guidance_scale, max_frames=None, frame_interval=1, preserve_frames=False, batch_size=8, progress=gr.Progress()):
277
  abort_event.clear() # Clear the abort flag at the start of a new job
278
  print("Starting video processing...")
279
+
280
 
281
  # Create a new job folder
282
  job_id = str(uuid.uuid4())