DiffMorpher/predict.py at main · edwardsanchez/DiffMorpher · GitHub

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
# Prediction interface for Cog ⚙️
# https://github.com/replicate/cog/blob/main/docs/python.md

import os
import shutil
from PIL import Image
import torch

from cog import BasePredictor, Input, Path
from model import DiffMorpherPipeline
from utils.lora_utils import train_lora


class Predictor(BasePredictor):
    def setup(self) -> None:
        """Load the model into memory to make running multiple predictions efficient"""
        self.model_name = "stabilityai/stable-diffusion-2-1-base"
        self.cache_dir = "model_cache"
        local_files_only = True  # set to True if the model is saved to cache_dir
        self.morpher_pipeline = DiffMorpherPipeline.from_pretrained(
            self.model_name,
            torch_dtype=torch.float32,
            cache_dir=self.cache_dir,
            local_files_only=local_files_only,
        ).to("cuda")

    def predict(
        self,
        image1: Path = Input(description="First input image."),
        image2: Path = Input(description="Second input image."),
        prompt1: str = Input(description="Description of the first input image."),
        prompt2: str = Input(description="Description of the second input image."),
        use_reschedule: bool = Input(
            description="Use reschedule sampling.", default=False
        ),
        use_adain: bool = Input(description="Use AdaIN.", default=False),
        lora_steps: int = Input(
            description="LoRA training steps.",
            default=200,
        ),
        lora_lr: float = Input(
            description="LoRA learning rate.",
            default=0.0002,
        ),
        lora_rank: int = Input(
            description="LoRA rank.",
            default=16,
        ),
        lamb: float = Input(
            description="Hyperparameter for self-attention replacement, where a larger indicates more replacements.",
            default=0.6,
        ),
        num_frames: int = Input(
            description="Number of frames to generate.", default=16
        ),
        duration: int = Input(
            description="Duration of each frame in millisecond.", default=100
        ),
    ) -> Path:
        """Run a single prediction on the model"""
        out_dir = "temp_dir"
        if os.path.exists(out_dir):
            shutil.rmtree(out_dir)
        os.makedirs(out_dir)

        lora1, lora2 = "lora_1.ckpt", "lora_2.ckpt"

        load_lora_path_1 = f"{out_dir}/{lora1}"
        load_lora_path_2 = f"{out_dir}/{lora2}"

        img1 = Image.open(str(image1)).convert("RGB")
        img2 = Image.open(str(image2)).convert("RGB")

        train_lora(
            img1,
            prompt1,
            out_dir,
            self.model_name,
            lora_steps=lora_steps,
            lora_lr=lora_lr,
            lora_rank=lora_rank,
            weight_name=lora1,
            cache_dir=self.cache_dir,
            local_files_only=True,
        )
        train_lora(
            img2,
            prompt2,
            out_dir,
            self.model_name,
            lora_steps=lora_steps,
            lora_lr=lora_lr,
            lora_rank=lora_rank,
            weight_name=lora2,
            cache_dir=self.cache_dir,
            local_files_only=True,
        )

        images = self.morpher_pipeline(
            img_path_0=str(image1),
            img_path_1=str(image2),
            prompt_0=prompt1,
            prompt_1=prompt2,
            save_lora_dir=out_dir,
            load_lora_path_0=load_lora_path_1,
            load_lora_path_1=load_lora_path_2,
            use_adain=use_adain,
            use_reschedule=use_reschedule,
            lamb=lamb,
            output_path=out_dir,
            num_frames=num_frames,
            fix_lora=None,
            save_intermediates=False,
        )
        out_path = "/tmp/out.gif"
        images[0].save(
            out_path, save_all=True, append_images=images[1:], duration=duration, loop=0
        )

        return Path(out_path)