Create deforum_engine.py

deforum_engine.py

@@ -0,0 +1,129 @@
+import torch, os, uuid, zipfile, cv2, gc, random
+import numpy as np
+from diffusers import AutoPipelineForImage2Image, LCMScheduler, EulerAncestralDiscreteScheduler, DDIMScheduler, DPMSolverMultistepScheduler
+from diffusers.models import AutoencoderTiny
+from PIL import Image
+import utils
+
+class DeforumRunner:
+    def __init__(self, device="cpu"):
+        self.device = device
+        self.pipe = None
+        self.stop_requested = False
+        self.current_config = (None, None, None)
+
+    def load_model(self, model_id, lora_id, scheduler_name):
+        if (model_id, lora_id, scheduler_name) == self.current_config and self.pipe: return
+        print(f"Loading: {model_id}...")
+        if self.pipe: del self.pipe; gc.collect()
+
+        # 1. Load Pipeline
+        try:
+            self.pipe = AutoPipelineForImage2Image.from_pretrained(
+                model_id, safety_checker=None, torch_dtype=torch.float32
+            )
+        except Exception as e:
+            print(f"Standard load failed: {e}. Trying explicit VAE fallback...")
+            # Fallback for SDXS or non-standard Repos
+            self.pipe = AutoPipelineForImage2Image.from_pretrained(
+                model_id, safety_checker=None, torch_dtype=torch.float32, use_safetensors=False
+            )
+
+        # 2. Handle LoRA
+        if lora_id and lora_id.lower() != "none":
+            try:
+                self.pipe.load_lora_weights(lora_id)
+                self.pipe.fuse_lora()
+                print(f"LoRA {lora_id} fused.")
+            except Exception as e: print(f"LoRA failed: {e}")
+
+        # 3. Handle Scheduler
+        config = self.pipe.scheduler.config
+        if scheduler_name == "LCM": self.pipe.scheduler = LCMScheduler.from_config(config)
+        elif scheduler_name == "Euler A": self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(config)
+        elif scheduler_name == "DDIM": self.pipe.scheduler = DDIMScheduler.from_config(config)
+        elif scheduler_name == "DPM++ 2M": self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(config)
+
+        self.pipe.to(self.device)
+        self.pipe.enable_attention_slicing()
+        self.current_config = (model_id, lora_id, scheduler_name)
+
+    def stop(self): self.stop_requested = True
+
+    def render(self, prompts, neg, max_f, w, h, z_s, a_s, tx_s, ty_s, str_s, noi_s, 
+               fps, steps, cfg, cadence, color_mode, border_mode, seed_beh, init_img, 
+               model_id, lora_id, scheduler_name):
+        
+        self.stop_requested = False
+        self.load_model(model_id, lora_id, scheduler_name)
+        
+        # Schedules
+        keys = ['z','a','tx','ty','str','noi']
+        vals = [z_s, a_s, tx_s, ty_s, str_s, noi_s]
+        sched = {k: utils.parse_weight_string(v, max_f) for k, v in zip(keys, vals)}
+        
+        run_id = uuid.uuid4().hex[:6]
+        os.makedirs(f"out_{run_id}", exist_ok=True)
+        
+        # Initial State
+        prev_img = init_img.resize((w, h), Image.LANCZOS) if init_img else Image.new("RGB", (w, h), (127, 127, 127))
+        color_anchor = prev_img.copy()
+        frames = []
+        
+        base_seed = random.randint(0, 2**32-1)
+        print(f"Run {run_id} | Model: {model_id} | Seed: {base_seed}")
+
+        for i in range(max_f):
+            if self.stop_requested: break
+            
+            # Seed Logic
+            f_seed = base_seed if seed_beh == "fixed" else (random.randint(0, 2**32-1) if seed_beh == "random" else base_seed + i)
+            random.seed(f_seed); np.random.seed(f_seed); torch.manual_seed(f_seed)
+            
+            # 1. WARP
+            args = {'angle':sched['a'][i], 'zoom':sched['z'][i], 'tx':sched['tx'][i], 'ty':sched['ty'][i]}
+            warped = utils.anim_frame_warp_2d(prev_img, args, border_mode)
+            
+            # 2. DIFFUSE (Cadence)
+            if i % cadence == 0:
+                # Color & Noise
+                inp = utils.maintain_colors(warped, color_anchor, color_mode)
+                inp = utils.add_noise(inp, sched['noi'][i])
+                
+                # Prompt
+                curr_prompt = prompts[max(k for k in prompts.keys() if k <= i)]
+                
+                # Strength (Ensure non-zero steps)
+                strength = sched['str'][i]
+                if steps * strength < 1.0: strength = 1.1 / steps
+                
+                # Generate
+                gen = self.pipe(
+                    prompt=curr_prompt, negative_prompt=neg, image=inp,
+                    num_inference_steps=steps, strength=strength, guidance_scale=cfg,
+                    width=w, height=h
+                ).images[0]
+                
+                if color_mode != 'None': gen = utils.maintain_colors(gen, color_anchor, color_mode)
+                prev_img = gen
+            else:
+                # Turbo
+                gen = warped
+                prev_img = warped
+                
+            frames.append(gen)
+            yield gen, None, None, f"Frame {i+1}/{max_f}"
+
+        # Save
+        vid_p = f"out_{run_id}/video.mp4"
+        out = cv2.VideoWriter(vid_p, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
+        for f in frames: out.write(cv2.cvtColor(np.array(f), cv2.COLOR_RGB2BGR))
+        out.release()
+        
+        zip_p = f"out_{run_id}/frames.zip"
+        with zipfile.ZipFile(zip_p, 'w') as zf:
+            for j, f in enumerate(frames):
+                buf = io.BytesIO(); f.save(buf, "PNG"); zf.writestr(f"{j:05d}.png", buf.getvalue())
+                import io # Re-import locally for safety
+        
+        yield frames[-1], vid_p, zip_p, "Done"