appiq-solution 1.2.0 → 1.3.0

package/#Tools/SkyReels-V2/skyreels_v2_infer/pipelines/image2video_pipeline.py

@@ -0,0 +1,156 @@
+import os
+from typing import List
+from typing import Optional
+from typing import Union
+
+import numpy as np
+import torch
+from diffusers.image_processor import PipelineImageInput
+from diffusers.video_processor import VideoProcessor
+from PIL import Image
+from tqdm import tqdm
+
+from ..modules import get_image_encoder
+from ..modules import get_text_encoder
+from ..modules import get_transformer
+from ..modules import get_vae
+from ..scheduler.fm_solvers_unipc import FlowUniPCMultistepScheduler
+
+
+def resizecrop(image: Image.Image, th, tw):
+    w, h = image.size
+    if w == tw and h == th:
+        return image
+    if h / w > th / tw:
+        new_w = int(w)
+        new_h = int(new_w * th / tw)
+    else:
+        new_h = int(h)
+        new_w = int(new_h * tw / th)
+    left = (w - new_w) / 2
+    top = (h - new_h) / 2
+    right = (w + new_w) / 2
+    bottom = (h + new_h) / 2
+    image = image.crop((left, top, right, bottom))
+    return image
+
+
+class Image2VideoPipeline:
+    def __init__(
+        self, model_path, dit_path, device: str = "cuda", weight_dtype=torch.bfloat16, use_usp=False, offload=False
+    ):
+        load_device = "cpu" if offload else device
+        self.transformer = get_transformer(dit_path, load_device, weight_dtype)
+        vae_model_path = os.path.join(model_path, "Wan2.1_VAE.pth")
+        self.vae = get_vae(vae_model_path, device, weight_dtype=torch.float32)
+        self.text_encoder = get_text_encoder(model_path, load_device, weight_dtype)
+        self.clip = get_image_encoder(model_path, load_device, weight_dtype)
+        self.sp_size = 1
+        self.device = device
+        self.offload = offload
+        self.video_processor = VideoProcessor(vae_scale_factor=16)
+        if use_usp:
+            from xfuser.core.distributed import get_sequence_parallel_world_size
+            from ..distributed.xdit_context_parallel import usp_attn_forward, usp_dit_forward
+            import types
+
+            for block in self.transformer.blocks:
+                block.self_attn.forward = types.MethodType(usp_attn_forward, block.self_attn)
+                self.transformer.forward = types.MethodType(usp_dit_forward, self.transformer)
+                self.sp_size = get_sequence_parallel_world_size()
+
+        self.scheduler = FlowUniPCMultistepScheduler()
+        self.vae_stride = (4, 8, 8)
+        self.patch_size = (1, 2, 2)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        image: PipelineImageInput,
+        prompt: Union[str, List[str]] = None,
+        negative_prompt: Union[str, List[str]] = None,
+        height: int = 544,
+        width: int = 960,
+        num_frames: int = 97,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 5.0,
+        shift: float = 5.0,
+        generator: Optional[torch.Generator] = None,
+    ):
+        F = num_frames
+
+        latent_height = height // 8 // 2 * 2
+        latent_width = width // 8 // 2 * 2
+        latent_length = (F - 1) // 4 + 1
+
+        h = latent_height * 8
+        w = latent_width * 8
+
+        img = self.video_processor.preprocess(image, height=h, width=w)
+
+        img = img.to(device=self.device, dtype=self.transformer.dtype)
+
+        padding_video = torch.zeros(img.shape[0], 3, F - 1, h, w, device=self.device)
+
+        img = img.unsqueeze(2)
+        img_cond = torch.concat([img, padding_video], dim=2)
+        img_cond = self.vae.encode(img_cond)
+        mask = torch.ones_like(img_cond)
+        mask[:, :, 1:] = 0
+        y = torch.cat([mask[:, :4], img_cond], dim=1)
+        self.clip.to(self.device)
+        clip_context = self.clip.encode_video(img)
+        if self.offload:
+            self.clip.cpu()
+            torch.cuda.empty_cache()
+
+        # preprocess
+        self.text_encoder.to(self.device)
+        context = self.text_encoder.encode(prompt).to(self.device)
+        context_null = self.text_encoder.encode(negative_prompt).to(self.device)
+        if self.offload:
+            self.text_encoder.cpu()
+            torch.cuda.empty_cache()
+
+        latent = torch.randn(
+            16, latent_length, latent_height, latent_width, dtype=torch.float32, generator=generator, device=self.device
+        )
+
+        self.transformer.to(self.device)
+        with torch.cuda.amp.autocast(dtype=self.transformer.dtype), torch.no_grad():
+            self.scheduler.set_timesteps(num_inference_steps, device=self.device, shift=shift)
+            timesteps = self.scheduler.timesteps
+
+            arg_c = {
+                "context": context,
+                "clip_fea": clip_context,
+                "y": y,
+            }
+
+            arg_null = {
+                "context": context_null,
+                "clip_fea": clip_context,
+                "y": y,
+            }
+
+            self.transformer.to(self.device)
+            for _, t in enumerate(tqdm(timesteps)):
+                latent_model_input = torch.stack([latent]).to(self.device)
+                timestep = torch.stack([t]).to(self.device)
+                noise_pred_cond = self.transformer(latent_model_input, t=timestep, **arg_c)[0].to(self.device)
+                noise_pred_uncond = self.transformer(latent_model_input, t=timestep, **arg_null)[0].to(self.device)
+                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+
+                temp_x0 = self.scheduler.step(
+                    noise_pred.unsqueeze(0), t, latent.unsqueeze(0), return_dict=False, generator=generator
+                )[0]
+                latent = temp_x0.squeeze(0)
+            if self.offload:
+                self.transformer.cpu()
+                torch.cuda.empty_cache()
+            videos = self.vae.decode(latent)
+            videos = (videos / 2 + 0.5).clamp(0, 1)
+            videos = [video for video in videos]
+            videos = [video.permute(1, 2, 3, 0) * 255 for video in videos]
+            videos = [video.cpu().numpy().astype(np.uint8) for video in videos]
+        return videos

package/#Tools/SkyReels-V2/skyreels_v2_infer/pipelines/prompt_enhancer.py

@@ -0,0 +1,65 @@
+import argparse
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+sys_prompt = """
+Transform the short prompt into a detailed video-generation caption using this structure:
+​​Opening shot type​​ (long/medium/close-up/extreme close-up/full shot)
+​​Primary subject(s)​​ with vivid attributes (colors, textures, actions, interactions)
+​​Dynamic elements​​ (movement, transitions, or changes over time, e.g., 'gradually lowers,' 'begins to climb,' 'camera moves toward...')
+​​Scene composition​​ (background, environment, spatial relationships)
+​​Lighting/atmosphere​​ (natural/artificial, time of day, mood)
+​​Camera motion​​ (zooms, pans, static/handheld shots) if applicable.
+
+Pattern Summary from Examples:
+[Shot Type] of [Subject+Action] + [Detailed Subject Description] + [Environmental Context] + [Lighting Conditions] + [Camera Movement]
+
+​One case: 
+Short prompt: a person is playing football
+Long prompt: Medium shot of a young athlete in a red jersey sprinting across a muddy field, dribbling a soccer ball with precise footwork. The player glances toward the goalpost, adjusts their stance, and kicks the ball forcefully into the net. Raindrops fall lightly, creating reflections under stadium floodlights. The camera follows the ball’s trajectory in a smooth pan.
+
+Note: If the subject is stationary, incorporate camera movement to ensure the generated video remains dynamic.
+
+​​Now expand this short prompt:​​ [{}]. Please only output the final long prompt in English.
+"""
+
+class PromptEnhancer:
+    def __init__(self, model_name="Qwen/Qwen2.5-32B-Instruct"):
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype="auto",
+            device_map="cuda:0",
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    def __call__(self, prompt):
+        prompt = prompt.strip()
+        prompt = sys_prompt.format(prompt)
+        messages = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": prompt}
+        ]
+        text = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        model_inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device)
+        generated_ids = self.model.generate(
+            **model_inputs,
+            max_new_tokens=2048,
+        )   
+        generated_ids = [
+            output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+        ]
+        rewritten_prompt = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        return rewritten_prompt
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--prompt", type=str, default="In a still frame, a stop sign")
+    args = parser.parse_args()
+
+    prompt_enhancer = PromptEnhancer()
+    enhanced_prompt = prompt_enhancer(args.prompt)
+    print(f'Original prompt: {args.prompt}')
+    print(f'Enhanced prompt: {enhanced_prompt}')

package/#Tools/SkyReels-V2/skyreels_v2_infer/pipelines/text2video_pipeline.py

@@ -0,0 +1,110 @@
+import os
+from typing import List
+from typing import Optional
+from typing import Union
+
+import numpy as np
+import torch
+from diffusers.video_processor import VideoProcessor
+from tqdm import tqdm
+
+from ..modules import get_text_encoder
+from ..modules import get_transformer
+from ..modules import get_vae
+from ..scheduler.fm_solvers_unipc import FlowUniPCMultistepScheduler
+
+
+class Text2VideoPipeline:
+    def __init__(
+        self, model_path, dit_path, device: str = "cuda", weight_dtype=torch.bfloat16, use_usp=False, offload=False
+    ):
+        load_device = "cpu" if offload else device
+        self.transformer = get_transformer(dit_path, load_device, weight_dtype)
+        vae_model_path = os.path.join(model_path, "Wan2.1_VAE.pth")
+        self.vae = get_vae(vae_model_path, device, weight_dtype=torch.float32)
+        self.text_encoder = get_text_encoder(model_path, load_device, weight_dtype)
+        self.video_processor = VideoProcessor(vae_scale_factor=16)
+        self.sp_size = 1
+        self.device = device
+        self.offload = offload
+        if use_usp:
+            from xfuser.core.distributed import get_sequence_parallel_world_size
+            from ..distributed.xdit_context_parallel import usp_attn_forward, usp_dit_forward
+            import types
+
+            for block in self.transformer.blocks:
+                block.self_attn.forward = types.MethodType(usp_attn_forward, block.self_attn)
+                self.transformer.forward = types.MethodType(usp_dit_forward, self.transformer)
+                self.sp_size = get_sequence_parallel_world_size()
+
+        self.scheduler = FlowUniPCMultistepScheduler()
+        self.vae_stride = (4, 8, 8)
+        self.patch_size = (1, 2, 2)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        negative_prompt: Union[str, List[str]] = None,
+        width: int = 544,
+        height: int = 960,
+        num_frames: int = 97,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 5.0,
+        shift: float = 5.0,
+        generator: Optional[torch.Generator] = None,
+    ):
+        # preprocess
+        F = num_frames
+        target_shape = (
+            self.vae.vae.z_dim,
+            (F - 1) // self.vae_stride[0] + 1,
+            height // self.vae_stride[1],
+            width // self.vae_stride[2],
+        )
+        self.text_encoder.to(self.device)
+        context = self.text_encoder.encode(prompt).to(self.device)
+        context_null = self.text_encoder.encode(negative_prompt).to(self.device)
+        if self.offload:
+            self.text_encoder.cpu()
+            torch.cuda.empty_cache()
+
+        latents = [
+            torch.randn(
+                target_shape[0],
+                target_shape[1],
+                target_shape[2],
+                target_shape[3],
+                dtype=torch.float32,
+                device=self.device,
+                generator=generator,
+            )
+        ]
+
+        # evaluation mode
+        self.transformer.to(self.device)
+        with torch.cuda.amp.autocast(dtype=self.transformer.dtype), torch.no_grad():
+            self.scheduler.set_timesteps(num_inference_steps, device=self.device, shift=shift)
+            timesteps = self.scheduler.timesteps
+
+            for _, t in enumerate(tqdm(timesteps)):
+                latent_model_input = torch.stack(latents)
+                timestep = torch.stack([t])
+                noise_pred_cond = self.transformer(latent_model_input, t=timestep, context=context)[0]
+                noise_pred_uncond = self.transformer(latent_model_input, t=timestep, context=context_null)[0]
+
+                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+
+                temp_x0 = self.scheduler.step(
+                    noise_pred.unsqueeze(0), t, latents[0].unsqueeze(0), return_dict=False, generator=generator
+                )[0]
+                latents = [temp_x0.squeeze(0)]
+            if self.offload:
+                self.transformer.cpu()
+                torch.cuda.empty_cache()
+            videos = self.vae.decode(latents[0])
+            videos = (videos / 2 + 0.5).clamp(0, 1)
+            videos = [video for video in videos]
+            videos = [video.permute(1, 2, 3, 0) * 255 for video in videos]
+            videos = [video.cpu().numpy().astype(np.uint8) for video in videos]
+        return videos