|
|
from typing import Any, Callable, Dict, List, Optional, Union |
|
|
|
|
|
import torch |
|
|
from diffusers import StableDiffusionImg2ImgPipeline |
|
|
from pipelines.twoStepPipeline import two_step_pipeline |
|
|
from util.commons import disable_safety_checker, download_image |
|
|
|
|
|
|
|
|
class Text2Img: |
|
|
def load(self, model_dir: str): |
|
|
self.pipe = two_step_pipeline.from_pretrained( |
|
|
model_dir, torch_dtype=torch.float16 |
|
|
).to("cuda") |
|
|
self.pipe.enable_xformers_memory_efficient_attention() |
|
|
disable_safety_checker(self.pipe) |
|
|
|
|
|
@torch.inference_mode() |
|
|
def process( |
|
|
self, |
|
|
prompt: Union[str, List[str]] = None, |
|
|
modified_prompts: Union[str, List[str]] = None, |
|
|
height: Optional[int] = None, |
|
|
width: Optional[int] = None, |
|
|
num_inference_steps: int = 50, |
|
|
guidance_scale: float = 7.5, |
|
|
negative_prompt: Optional[Union[str, List[str]]] = None, |
|
|
num_images_per_prompt: Optional[int] = 1, |
|
|
eta: float = 0.0, |
|
|
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, |
|
|
latents: Optional[torch.FloatTensor] = None, |
|
|
prompt_embeds: Optional[torch.FloatTensor] = None, |
|
|
negative_prompt_embeds: Optional[torch.FloatTensor] = None, |
|
|
output_type: Optional[str] = "pil", |
|
|
return_dict: bool = True, |
|
|
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, |
|
|
callback_steps: int = 1, |
|
|
cross_attention_kwargs: Optional[Dict[str, Any]] = None, |
|
|
iteration: float = 3.0, |
|
|
): |
|
|
return self.pipe.two_step_pipeline( |
|
|
prompt=prompt, |
|
|
modified_prompts=modified_prompts, |
|
|
height=height, |
|
|
width=width, |
|
|
num_inference_steps=num_inference_steps, |
|
|
guidance_scale=guidance_scale, |
|
|
negative_prompt=negative_prompt, |
|
|
num_images_per_prompt=num_images_per_prompt, |
|
|
eta=eta, |
|
|
generator=generator, |
|
|
latents=latents, |
|
|
prompt_embeds=prompt_embeds, |
|
|
negative_prompt_embeds=negative_prompt_embeds, |
|
|
output_type=output_type, |
|
|
return_dict=return_dict, |
|
|
callback=callback, |
|
|
callback_steps=callback_steps, |
|
|
cross_attention_kwargs=cross_attention_kwargs, |
|
|
iteration=iteration, |
|
|
).images |
|
|
|
|
|
|
|
|
class Img2Img: |
|
|
def load(self, model_dir: str): |
|
|
self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained( |
|
|
model_dir, torch_dtype=torch.float16 |
|
|
).to("cuda") |
|
|
self.pipe.enable_xformers_memory_efficient_attention() |
|
|
disable_safety_checker(self.pipe) |
|
|
|
|
|
@torch.inference_mode() |
|
|
def process( |
|
|
self, prompt: List[str], imageUrl: str, negative_prompt: List[str], steps: int |
|
|
): |
|
|
image = download_image(imageUrl) |
|
|
|
|
|
return self.pipe.__call__( |
|
|
prompt=prompt, |
|
|
image=image, |
|
|
strength=0.75, |
|
|
negative_prompt=negative_prompt, |
|
|
guidance_scale=7.5, |
|
|
num_images_per_prompt=1, |
|
|
num_inference_steps=steps, |
|
|
).images |
|
|
|
