File size: 1,908 Bytes
413d4d0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
class CogVideo:
    def __init__(self, device="cuda"):
        """
        Initializes the CogVideo model with a specific device.

        Args:
            device (str, optional): The device to run the model on. Defaults to "cuda".
        """

        import argparse

        # Manually creating an args object
        self.args = argparse.Namespace(
            generate_frame_num=5,
            coglm_temperature2=0.89,
            use_guidance_stage1=True,
            use_guidance_stage2=False,  # Assuming this is not set
            guidance_alpha=3.0,
            stage_1=False,  # Assuming this is not set
            stage_2=False,  # Assuming this is not set
            both_stages=True,
            parallel_size=1,
            stage1_max_inference_batch_size=-1,
            multi_gpu=False,  # Assuming this is not set
            device=3,
        )

    def infer_one_video(
        self,
        prompt: str = None,
        size: list = [320, 512],
        seconds: int = 2,
        fps: int = 8,
        seed: int = 42,
    ):
        """
        Generates a single video based on the provided prompt and parameters.

        Args:
            prompt (str, optional): The text prompt to generate the video from. Defaults to None.
            size (list, optional): The size of the video as [height, width]. Defaults to [320, 512].
            seconds (int, optional): The duration of the video in seconds. Defaults to 2.
            fps (int, optional): The frames per second of the video. Defaults to 8.
            seed (int, optional): The seed for random number generation. Defaults to 42.

        Returns:
            torch.Tensor: The generated video as a tensor.
        """

        from videogen_hub.pipelines.cogvideo.cogvideo_pipeline import pipeline

        return pipeline(
            self.args, raw_text=prompt, height=size[0], width=size[1], duration=seconds
        )