File size: 3,117 Bytes
4adca93 |
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 |
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
|