|
import math |
|
from typing import List, Optional |
|
|
|
from PIL import Image |
|
|
|
from internals.data.result import Result |
|
from internals.pipelines.commons import AbstractPipeline, Img2Img |
|
from internals.util.config import get_model_dir, get_base_dimension |
|
|
|
|
|
class HighRes(AbstractPipeline): |
|
def load(self, img2img: Optional[Img2Img] = None): |
|
if hasattr(self, "pipe"): |
|
return |
|
|
|
if not img2img: |
|
img2img = Img2Img() |
|
img2img.load(get_model_dir()) |
|
|
|
self.pipe = img2img.pipe |
|
self.img2img = img2img |
|
|
|
def apply( |
|
self, |
|
prompt: List[str], |
|
negative_prompt: List[str], |
|
images, |
|
width: int, |
|
height: int, |
|
num_inference_steps: int, |
|
strength: float = 0.5, |
|
guidance_scale: int = 9, |
|
**kwargs, |
|
): |
|
images = [image.resize((width, height)) for image in images] |
|
kwargs = { |
|
"prompt": prompt, |
|
"image": images, |
|
"strength": strength, |
|
"negative_prompt": negative_prompt, |
|
"guidance_scale": guidance_scale, |
|
"num_inference_steps": num_inference_steps, |
|
**kwargs, |
|
} |
|
result = self.pipe.__call__(**kwargs) |
|
return Result.from_result(result) |
|
|
|
@staticmethod |
|
def get_intermediate_dimension(target_width: int, target_height: int): |
|
def_size = get_base_dimension() |
|
|
|
desired_pixel_count = def_size * def_size |
|
actual_pixel_count = target_width * target_height |
|
|
|
scale = math.sqrt(desired_pixel_count / actual_pixel_count) |
|
|
|
firstpass_width = math.ceil(scale * target_width / 64) * 64 |
|
firstpass_height = math.ceil(scale * target_height / 64) * 64 |
|
|
|
return firstpass_width, firstpass_height |
|
|
