Create deforum_engine.py

deforum_engine.py

@@ -0,0 +1,154 @@
+import torch
+import numpy as np
+from diffusers import AutoPipelineForImage2Image, LCMScheduler, EulerAncestralDiscreteScheduler
+from PIL import Image
+import utils
+import os
+import uuid
+import zipfile
+import cv2
+import gc
+
+class DeforumRunner:
+    def __init__(self, device="cpu"):
+        self.device = device
+        self.pipe = None
+        self.stop_requested = False
+        self.current_model_config = (None, None, None)
+
+    def load_model(self, model_id, lora_id, scheduler_name):
+        """Loads model, LoRA, and scheduler dynamically."""
+        new_config = (model_id, lora_id, scheduler_name)
+        if new_config == self.current_model_config and self.pipe is not None:
+            return
+
+        print(f"Loading Model: {model_id}, LoRA: {lora_id}, Scheduler: {scheduler_name}")
+        if self.pipe:
+            del self.pipe
+            gc.collect()
+
+        try:
+            pipe = AutoPipelineForImage2Image.from_pretrained(
+                model_id, safety_checker=None, torch_dtype=torch.float32
+            )
+        except Exception as e:
+            print(f"Error loading model {model_id}: {e}. Falling back.")
+            pipe = AutoPipelineForImage2Image.from_pretrained(
+                "runwayml/stable-diffusion-v1-5", safety_checker=None, torch_dtype=torch.float32
+            )
+
+        if lora_id and lora_id != "None":
+            try:
+                pipe.load_lora_weights(lora_id)
+                pipe.fuse_lora()
+                print("LoRA loaded and fused.")
+            except Exception as e: print(f"Error loading LoRA: {e}")
+
+        if scheduler_name == "LCM":
+            pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+        elif scheduler_name == "Euler A":
+            pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+
+        pipe.to(self.device)
+        pipe.set_progress_bar_config(disable=True)
+        pipe.enable_attention_slicing() # Crucial for CPU memory
+        self.pipe = pipe
+        self.current_model_config = new_config
+        print("Pipeline ready.")
+
+    def stop(self):
+        self.stop_requested = True
+
+    def render(self, prompts, neg_prompt, max_frames, width, height,
+               zoom_s, angle_s, tx_s, ty_s, strength_s, noise_s, 
+               fps, steps, cadence, color_mode, border_mode, init_image_upload,
+               model_id, lora_id, scheduler_name):
+        
+        self.stop_requested = False
+        self.load_model(model_id, lora_id, scheduler_name)
+        
+        # 1. Parse Schedules
+        keys = ['z', 'a', 'tx', 'ty', 'str', 'noi']
+        inputs = [zoom_s, angle_s, tx_s, ty_s, strength_s, noise_s]
+        sched = {k: utils.parse_weight_string(v, max_frames) for k, v in zip(keys, inputs)}
+        
+        # 2. Setup
+        run_id = uuid.uuid4().hex[:6]
+        output_dir = f"output_{run_id}"
+        os.makedirs(output_dir, exist_ok=True)
+        
+        if init_image_upload:
+            prev_img = init_image_upload.resize((width, height), Image.LANCZOS)
+            color_anchor = prev_img
+        else:
+            prev_img = None
+            color_anchor = None
+            
+        generated_frames = []
+        print(f"Starting run {run_id}...")
+
+        # 3. Loop
+        for i in range(max_frames):
+            if self.stop_requested:
+                print("Generation stopped.")
+                break
+            
+            # Get Params
+            z, a, tx, ty = sched['z'][i], sched['a'][i], sched['tx'][i], sched['ty'][i]
+            strength, noise = sched['str'][i], sched['noi'][i]
+            current_prompt = prompts[max(k for k in prompts.keys() if k <= i)]
+            
+            # --- Authentic Deforum Loop ---
+            # 1. Warp Previous Frame
+            if prev_img is not None:
+                warped_img = utils.anim_frame_warp_2d(prev_img, {'angle': a, 'zoom': z, 'translation_x': tx, 'translation_y': ty}, border_mode)
+            else:
+                warped_img = Image.new("RGB", (width, height), (0,0,0))
+
+            # Decide: Diffusion or Just Warp (Cadence)
+            if i % cadence == 0:
+                # 2. Color Match & 3. Add Noise (Only before diffusion)
+                init_for_diff = utils.maintain_colors(warped_img, color_anchor, color_mode)
+                init_for_diff = utils.add_noise(init_for_diff, noise)
+                
+                # 4. Diffusion (Img2Img)
+                # Use high strength for first frame if no init image provided
+                curr_strength = strength if prev_img is not None else 0.95
+                
+                gen_image = self.pipe(
+                    prompt=current_prompt, negative_prompt=neg_prompt,
+                    image=init_for_diff, num_inference_steps=steps,
+                    strength=curr_strength, guidance_scale=1.2, # Low CFG for LCM
+                    width=width, height=height
+                ).images[0]
+            else:
+                # Cadence step: Just show the warped image (faster)
+                gen_image = warped_img
+
+            # Update state
+            prev_img = gen_image
+            if color_anchor is None: color_anchor = gen_image
+            generated_frames.append(gen_image)
+            yield gen_image, None, None
+
+        # 4. Finalize
+        vid_path = f"{output_dir}/video.mp4"
+        self.save_video(generated_frames, vid_path, fps)
+        zip_path = f"{output_dir}/frames.zip"
+        self.save_zip(generated_frames, zip_path)
+        yield generated_frames[-1], vid_path, zip_path
+
+    # (save_video and save_zip are the same as before, omitted for brevity)
+    def save_video(self, frames, path, fps):
+        if not frames: return
+        w, h = frames[0].size
+        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))
+        out.release()
+    def save_zip(self, frames, path):
+        import io
+        with zipfile.ZipFile(path, 'w') as zf:
+            for i, f in enumerate(frames):
+                buf = io.BytesIO()
+                f.save(buf, format="PNG")
+                zf.writestr(f"{i:05d}.png", buf.getvalue())