Update deforum_engine.py

deforum_engine.py

@@ -12,40 +12,44 @@ class DeforumRunner:
         self.current_model_config = (None, None, None) # Model, LoRA, Scheduler
 
     def load_model(self, model_id, lora_id, scheduler_name):
-        # Prevent reloading if config matches
+        # Prevent reloading if config matches exactly
         if (model_id, lora_id, scheduler_name) == self.current_model_config and self.pipe is not None:
             return
 
         print(f"Loading Model: {model_id}")
         
-        # Safe Cleanup
-        if hasattr(self, 'pipe') and self.pipe is not None:
+        # Safe Cleanup: Check existence before deletion to prevent AttributeError
+        if getattr(self, 'pipe', None) is not None:
             del self.pipe
         self.pipe = None
         gc.collect()
 
-        # --- STRATEGY 1: Standard Load ---
-        try:
-            self.pipe = AutoPipelineForImage2Image.from_pretrained(
-                model_id, safety_checker=None, torch_dtype=torch.float32
-            )
-        except Exception as e1:
-            print(f"Strategy 1 failed ({e1}). Trying Legacy (BIN)...")
-            # --- STRATEGY 2: Legacy .bin files (use_safetensors=False) ---
+        # Loading Strategies: Try them in order until one works
+        strategies = [
+            # 1. Standard Load (Safetensors default)
+            {"safety_checker": None, "torch_dtype": torch.float32},
+            # 2. Legacy Load (PyTorch .bin files)
+            {"safety_checker": None, "torch_dtype": torch.float32, "use_safetensors": False},
+            # 3. FP16 Variant (Common in optimized repos, but we cast to float32 for CPU safety)
+            {"safety_checker": None, "torch_dtype": torch.float32, "variant": "fp16"},
+        ]
+
+        success = False
+        last_error = None
+
+        for kwargs in strategies:
             try:
-                self.pipe = AutoPipelineForImage2Image.from_pretrained(
-                    model_id, safety_checker=None, torch_dtype=torch.float32, use_safetensors=False
-                )
-            except Exception as e2:
-                print(f"Strategy 2 failed ({e2}). Trying FP16 Variant...")
-                # --- STRATEGY 3: FP16 Variant ---
-                try:
-                    self.pipe = AutoPipelineForImage2Image.from_pretrained(
-                        model_id, safety_checker=None, torch_dtype=torch.float32, variant="fp16"
-                    )
-                except Exception as e3:
-                    print(f"All loading strategies failed for {model_id}.")
-                    raise RuntimeError(f"Could not load model {model_id}. Last error: {e3}")
+                self.pipe = AutoPipelineForImage2Image.from_pretrained(model_id, **kwargs)
+                success = True
+                print(f"Loaded successfully with config: {kwargs}")
+                break
+            except Exception as e:
+                last_error = e
+                continue
+
+        if not success:
+            print(f"All loading strategies failed for {model_id}.")
+            raise RuntimeError(f"Could not load model {model_id}. Error: {last_error}")
 
         # Load LoRA
         if lora_id and lora_id.strip().lower() != "none":
@@ -53,7 +57,8 @@ class DeforumRunner:
                 self.pipe.load_lora_weights(lora_id)
                 self.pipe.fuse_lora()
                 print(f"LoRA {lora_id} loaded.")
-            except Exception as e: print(f"LoRA Load Error: {e}")
+            except Exception as e: 
+                print(f"LoRA Load Error (continuing without LoRA): {e}")
 
         # Configure Scheduler
         s_config = self.pipe.scheduler.config
@@ -67,7 +72,8 @@ class DeforumRunner:
             self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(s_config)
 
         self.pipe.to(self.device)
-        # Check if attention slicing is supported (TinyVAE sometimes hates it)
+        
+        # TinyVAE / specific models crash with attention slicing, so we wrap it
         try: self.pipe.enable_attention_slicing()
         except: pass
         
@@ -85,13 +91,15 @@ class DeforumRunner:
                model_path, lora_path, scheduler_type):
         
         self.stop_requested = False
+        
+        # Load Model (Propagate errors to UI)
         try:
             self.load_model(model_path, lora_path, scheduler_type)
         except Exception as e:
-            yield None, None, None, f"Model Load Error: {str(e)}"
+            yield None, None, None, f"Model Error: {str(e)}"
             return
         
-        # 1. Parse Schedules
+        # Parse Schedules
         try:
             schedules = {
                 'zoom': utils.parse_weight_schedule(zoom_schedule, max_frames),
@@ -108,7 +116,7 @@ class DeforumRunner:
         run_id = uuid.uuid4().hex[:6]
         os.makedirs(f"output_{run_id}", exist_ok=True)
         
-        # 2. Init Canvas
+        # Init Canvas
         if init_image:
             prev_img = init_image.resize((width, height), Image.LANCZOS)
         else:
@@ -120,7 +128,7 @@ class DeforumRunner:
         base_seed = random.randint(0, 2**32 - 1)
         print(f"Run {run_id} Started. Base Seed: {base_seed}")
 
-        # 3. Main Loop
+        # Main Loop
         for frame_idx in range(max_frames):
             if self.stop_requested: break
             
@@ -133,7 +141,7 @@ class DeforumRunner:
             np.random.seed(current_seed)
             torch.manual_seed(current_seed)
             
-            # Warp
+            # 1. Warp
             warp_args = {
                 'angle': schedules['angle'][frame_idx], 
                 'zoom': schedules['zoom'][frame_idx], 
@@ -142,7 +150,7 @@ class DeforumRunner:
             }
             warped_img = utils.anim_frame_warp_2d(prev_img, warp_args, border_mode)
             
-            # Cadence Logic
+            # 2. Cadence Logic
             if frame_idx % cadence == 0:
                 # Prepare
                 init_for_diffusion = utils.maintain_colors(warped_img, color_anchor, color_coherence)
@@ -171,16 +179,16 @@ class DeforumRunner:
                 
                 prev_img = gen_image
             else:
-                # Turbo
+                # Turbo (Warp Only)
                 gen_image = warped_img
                 prev_img = warped_img
 
             generated_frames.append(gen_image)
             
-            # Yield Frame
+            # Yield Status
             yield gen_image, None, None, f"Frame {frame_idx+1}/{max_frames}"
 
-        # 4. Save Outputs
+        # Finalize
         video_path = f"output_{run_id}/video.mp4"
         self.save_video(generated_frames, video_path, fps)
         
@@ -193,7 +201,7 @@ class DeforumRunner:
         if not frames: return
         try:
             w, h = frames[0].size
-            # Try 'mp4v' first
+            # 'mp4v' is widely supported for CPU/OpenCV
             out = cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
             for f in frames:
                 out.write(cv2.cvtColor(np.array(f), cv2.COLOR_RGB2BGR))