Create deforum_engine.py

deforum_engine.py

@@ -0,0 +1,150 @@
+import torch
+import numpy as np
+from diffusers import StableDiffusionImg2ImgPipeline, EulerAncestralDiscreteScheduler
+from diffusers.utils import load_image
+import utils
+import os
+import uuid
+import zipfile
+import cv2
+
+class DeforumRunner:
+    def __init__(self, device="cpu"):
+        self.device = device
+        # Using the requested model
+        # Note: We use Img2ImgPipeline. Even for SDXS, we need 'steps' to make 'strength' work.
+        self.model_id = "IDKiro/sdxs-512-dreamshaper"
+        self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+            self.model_id, 
+            safety_checker=None,
+            torch_dtype=torch.float32
+        )
+        self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config)
+        self.pipe.to(self.device)
+        self.pipe.set_progress_bar_config(disable=True)
+
+    def render(self, 
+               prompts,         # Dict {0: "prompt"}
+               neg_prompt,
+               max_frames,
+               width, height,
+               zoom_str, angle_str, tx_str, ty_str,
+               strength_str, noise_str,
+               fps, steps):
+        
+        # 1. Parse all schedules up front
+        schedules = {
+            'zoom': utils.parse_weight_string(zoom_str, max_frames),
+            'angle': utils.parse_weight_string(angle_str, max_frames),
+            'tx': utils.parse_weight_string(tx_str, max_frames),
+            'ty': utils.parse_weight_string(ty_str, max_frames),
+            'strength': utils.parse_weight_string(strength_str, max_frames),
+            'noise': utils.parse_weight_string(noise_str, max_frames)
+        }
+        
+        # 2. Setup output
+        run_id = uuid.uuid4().hex[:6]
+        output_dir = f"output_{run_id}"
+        os.makedirs(output_dir, exist_ok=True)
+        
+        prev_img = None
+        color_anchor = None # For 'Match Frame 0'
+        
+        generated_frames = []
+        
+        print(f"Starting Deforum Run {run_id} for {max_frames} frames...")
+
+        # 3. Main Loop
+        for frame_idx in range(max_frames):
+            
+            # --- A. Get Parameters for this frame ---
+            zoom = schedules['zoom'][frame_idx]
+            angle = schedules['angle'][frame_idx]
+            tx = schedules['tx'][frame_idx]
+            ty = schedules['ty'][frame_idx]
+            noise_amt = schedules['noise'][frame_idx]
+            strength = schedules['strength'][frame_idx]
+            
+            # Get prompt (find latest key <= current frame)
+            prompt_keys = sorted([k for k in prompts.keys() if k <= frame_idx])
+            current_prompt = prompts[prompt_keys[-1]]
+            
+            # --- B. Prepare Init Image ---
+            if prev_img is None:
+                # First frame: Generate from scratch (Text2Img via Img2Img with strength=1 effectively)
+                # We use a blank image or random noise as base if strictly using Img2Img, 
+                # but better to just use strength=1.0 or high strength.
+                init_image = Image.new("RGB", (width, height), (0,0,0)) # Dummy
+                # High strength tells pipe to ignore init image
+                current_strength = 1.0 
+            else:
+                # WARPING logic
+                warp_args = {'angle': angle, 'zoom': zoom, 'translation_x': tx, 'translation_y': ty}
+                init_image = utils.anim_frame_warp_2d(prev_img, warp_args)
+                
+                # COLOR MATCHING (Match Frame 0)
+                if color_anchor is not None:
+                    init_image = utils.maintain_colors(init_image, color_anchor)
+                
+                # NOISE INJECTION
+                init_image = utils.add_noise(init_image, noise_amt)
+                
+                current_strength = strength
+
+            # --- C. Generation ---
+            # If we are strictly 1-step, strength is ignored. 
+            # We enforce min 2 steps if strength < 1.0 to ensure mixing.
+            # However, user wants SDXS optimization. 
+            # Compromise: We pass the standard strength. 
+            # If steps=1, strength is usually moot. We rely on 'steps' from UI (recommend 3).
+            
+            gen_image = self.pipe(
+                prompt=current_prompt,
+                negative_prompt=neg_prompt,
+                image=init_image,
+                num_inference_steps=int(steps), 
+                strength=current_strength,
+                guidance_scale=0.0, # SDXS Requirement
+                width=width,
+                height=height
+            ).images[0]
+            
+            # Update state
+            prev_img = gen_image
+            if color_anchor is None:
+                color_anchor = gen_image
+            
+            generated_frames.append(gen_image)
+            
+            # Yield results for UI
+            yield gen_image, None, None
+            
+        # 4. Finalize (Video & Zip)
+        video_path = os.path.join(output_dir, "video.mp4")
+        self.save_video(generated_frames, video_path, fps)
+        
+        zip_path = os.path.join(output_dir, "frames.zip")
+        self.save_zip(generated_frames, zip_path)
+        
+        yield generated_frames[-1], video_path, zip_path
+
+    def save_video(self, frames, path, fps):
+        if not frames: return
+        w, h = frames[0].size
+        # Use OpenCV for CPU efficiency
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        out = cv2.VideoWriter(path, fourcc, fps, (w, h))
+        for f in frames:
+            # PIL RGB -> OpenCV BGR
+            out.write(cv2.cvtColor(np.array(f), cv2.COLOR_RGB2BGR))
+        out.release()
+
+    def save_zip(self, frames, path):
+        with zipfile.ZipFile(path, 'w') as zf:
+            for i, f in enumerate(frames):
+                name = f"{i:05d}.png"
+                # Save to buffer to avoid writing temp files to disk
+                import io
+                buf = io.BytesIO()
+                f.save(buf, format="PNG")
+                zf.writestr(name, buf.getvalue())