File size: 5,136 Bytes
eb65e9e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
import os

import gradio as gr
import torch
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
from PIL import Image

from video_diffusion.inpaint_zoom.utils.zoom_out_utils import (
    dummy,
    preprocess_image,
    preprocess_mask_image,
    write_video,
)

os.environ["CUDA_VISIBLE_DEVICES"] = "0"


stable_paint_model_list = ["stabilityai/stable-diffusion-2-inpainting", "runwayml/stable-diffusion-inpainting"]

stable_paint_prompt_list = [
    "children running in the forest , sunny, bright, by studio ghibli painting, superior quality, masterpiece,  traditional Japanese colors, by Grzegorz Rutkowski, concept art",
    "A beautiful landscape of a mountain range with a lake in the foreground",
]

stable_paint_negative_prompt_list = [
    "lurry, bad art, blurred, text, watermark",
]


class StableDiffusionZoomOut:
    def __init__(self):
        self.pipe = None

    def load_model(self, model_id):
        if self.pipe is None:
            self.pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
            self.pipe.set_use_memory_efficient_attention_xformers(True)
            self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
            self.pipe = self.pipe.to("cuda")
            self.pipe.safety_checker = dummy
            self.g_cuda = torch.Generator(device="cuda")

        return self.pipe

    def generate_video(
        self,
        model_id,
        prompt,
        negative_prompt,
        guidance_scale,
        num_inference_steps,
        num_frames,
        step_size,
    ):
        pipe = self.load_model(model_id)

        new_image = Image.new(mode="RGBA", size=(512, 512))
        current_image, mask_image = preprocess_mask_image(new_image)

        current_image = pipe(
            prompt=[prompt],
            negative_prompt=[negative_prompt],
            image=current_image,
            mask_image=mask_image,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
        ).images[0]

        all_frames = []
        all_frames.append(current_image)

        for i in range(num_frames):
            prev_image = preprocess_image(current_image, step_size, 512)
            current_image = prev_image
            current_image, mask_image = preprocess_mask_image(current_image)
            current_image = pipe(
                prompt=[prompt],
                negative_prompt=[negative_prompt],
                image=current_image,
                mask_image=mask_image,
                num_inference_steps=num_inference_steps,
            ).images[0]
            current_image.paste(prev_image, mask=prev_image)
            all_frames.append(current_image)

        save_path = "output.mp4"
        write_video(save_path, all_frames, fps=30)
        return save_path

    def app():
        with gr.Blocks():
            with gr.Row():
                with gr.Column():
                    text2image_out_model_path = gr.Dropdown(
                        choices=stable_paint_model_list, value=stable_paint_model_list[0], label="Text-Image Model Id"
                    )

                    text2image_out_prompt = gr.Textbox(lines=2, value=stable_paint_prompt_list[0], label="Prompt")

                    text2image_out_negative_prompt = gr.Textbox(
                        lines=1, value=stable_paint_negative_prompt_list[0], label="Negative Prompt"
                    )

                    with gr.Row():
                        with gr.Column():
                            text2image_out_guidance_scale = gr.Slider(
                                minimum=0.1, maximum=15, step=0.1, value=7.5, label="Guidance Scale"
                            )

                            text2image_out_num_inference_step = gr.Slider(
                                minimum=1, maximum=100, step=1, value=50, label="Num Inference Step"
                            )
                        with gr.Row():
                            with gr.Column():
                                text2image_out_step_size = gr.Slider(
                                    minimum=1, maximum=100, step=1, value=10, label="Step Size"
                                )

                                text2image_out_num_frames = gr.Slider(
                                    minimum=1, maximum=100, step=1, value=10, label="Frames"
                                )

                    text2image_out_predict = gr.Button(value="Generator")

                with gr.Column():
                    output_image = gr.Video(label="Output")

            text2image_out_predict.click(
                fn=StableDiffusionZoomOut().generate_video,
                inputs=[
                    text2image_out_model_path,
                    text2image_out_prompt,
                    text2image_out_negative_prompt,
                    text2image_out_guidance_scale,
                    text2image_out_num_inference_step,
                    text2image_out_step_size,
                    text2image_out_num_frames,
                ],
                outputs=output_image,
            )