Update lollms_fooocus.py

This commit is contained in:
Saifeddine ALOUI 2024-05-31 11:09:29 +02:00 committed by GitHub
parent b6fd8d0434
commit 942fced394
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -54,14 +54,14 @@ def download_file(url, folder_path, local_filename):
def install_model(lollms_app:LollmsApplication, model_url): def install_model(lollms_app:LollmsApplication, model_url):
root_dir = lollms_app.lollms_paths.personal_path root_dir = lollms_app.lollms_paths.personal_path
shared_folder = root_dir/"shared" shared_folder = root_dir/"shared"
diffusers_folder = shared_folder / "diffusers" fooocus_folder = shared_folder / "fooocus"
if not PackageManager.check_package_installed("diffusers"): if not PackageManager.check_package_installed("fooocus"):
PackageManager.install_or_update("diffusers") PackageManager.install_or_update("fooocus")
if not PackageManager.check_package_installed("torch"): if not PackageManager.check_package_installed("torch"):
check_and_install_torch(True) check_and_install_torch(True)
import torch import torch
from diffusers import PixArtSigmaPipeline from fooocus import PixArtSigmaPipeline
# You can replace the checkpoint id with "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" too. # You can replace the checkpoint id with "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" too.
pipe = PixArtSigmaPipeline.from_pretrained( pipe = PixArtSigmaPipeline.from_pretrained(
@ -69,32 +69,32 @@ def install_model(lollms_app:LollmsApplication, model_url):
) )
def install_diffusers(lollms_app:LollmsApplication): def install_fooocus(lollms_app:LollmsApplication):
root_dir = lollms_app.lollms_paths.personal_path root_dir = lollms_app.lollms_paths.personal_path
shared_folder = root_dir/"shared" shared_folder = root_dir/"shared"
diffusers_folder = shared_folder / "diffusers" fooocus_folder = shared_folder / "fooocus"
diffusers_folder.mkdir(exist_ok=True, parents=True) fooocus_folder.mkdir(exist_ok=True, parents=True)
if not PackageManager.check_package_installed("diffusers"): if not PackageManager.check_package_installed("fooocus"):
PackageManager.install_or_update("diffusers") PackageManager.install_or_update("gradio_client")
PackageManager.install_or_update("xformers")
def upgrade_diffusers(lollms_app:LollmsApplication): def upgrade_fooocus(lollms_app:LollmsApplication):
PackageManager.install_or_update("diffusers") PackageManager.install_or_update("fooocus")
PackageManager.install_or_update("xformers") PackageManager.install_or_update("xformers")
class LollmsFooocus(LollmsTTI): class LollmsFooocus(LollmsTTI):
has_controlnet = False
def __init__( def __init__(
self, self,
app:LollmsApplication, app:LollmsApplication,
wm = "Artbot", wm = "Artbot",
base_url="localhost:1024"
): ):
super().__init__(app) super().__init__(app)
self.ready = False self.ready = False
self.base_url = base_url
# Get the current directory # Get the current directory
lollms_paths = app.lollms_paths lollms_paths = app.lollms_paths
root_dir = lollms_paths.personal_path root_dir = lollms_paths.personal_path
@ -102,9 +102,9 @@ class LollmsFooocus(LollmsTTI):
self.wm = wm self.wm = wm
shared_folder = root_dir/"shared" shared_folder = root_dir/"shared"
self.diffusers_folder = shared_folder / "diffusers" self.fooocus_folder = shared_folder / "fooocus"
self.output_dir = root_dir / "outputs/diffusers" self.output_dir = root_dir / "outputs/fooocus"
self.models_dir = self.diffusers_folder / "models" self.models_dir = self.fooocus_folder / "models"
self.output_dir.mkdir(parents=True, exist_ok=True) self.output_dir.mkdir(parents=True, exist_ok=True)
self.models_dir.mkdir(parents=True, exist_ok=True) self.models_dir.mkdir(parents=True, exist_ok=True)
@ -116,38 +116,32 @@ class LollmsFooocus(LollmsTTI):
ASCIIColors.red("\_____/\___/|_|_|_| |_| |_|___/ \_| \___/ \___/ \___/ \___|\__,_|___/") ASCIIColors.red("\_____/\___/|_|_|_| |_| |_|___/ \_| \___/ \___/ \___/ \___|\__,_|___/")
ASCIIColors.red(" ______ ") ASCIIColors.red(" ______ ")
ASCIIColors.red(" |______| ") ASCIIColors.red(" |______| ")
import torch if not PackageManager.check_package_installed("gradio_client"):
from diffusers import PixArtSigmaPipeline PackageManager.install_or_update("gradio_client")
self.model = PixArtSigmaPipeline.from_pretrained( from gradio_client import Client
app.config.diffusers_model, torch_dtype=torch.float16, cache_dir=self.models_dir, self.client = Client(base_url)
use_safetensors=True,
)
# Enable memory optimizations.
if app.config.diffusers_offloading_mode=="sequential_cpu_offload":
self.model.enable_sequential_cpu_offload()
elif app.coinfig.diffusers_offloading_mode=="model_cpu_offload":
self.model.enable_model_cpu_offload()
@staticmethod @staticmethod
def verify(app:LollmsApplication): def verify(app:LollmsApplication):
# Clone repository # Clone repository
root_dir = app.lollms_paths.personal_path root_dir = app.lollms_paths.personal_path
shared_folder = root_dir/"shared" shared_folder = root_dir/"shared"
diffusers_folder = shared_folder / "diffusers" fooocus_folder = shared_folder / "fooocus"
return diffusers_folder.exists() return fooocus_folder.exists()
def get(app:LollmsApplication): def get(app:LollmsApplication):
root_dir = app.lollms_paths.personal_path root_dir = app.lollms_paths.personal_path
shared_folder = root_dir/"shared" shared_folder = root_dir/"shared"
diffusers_folder = shared_folder / "diffusers" fooocus_folder = shared_folder / "fooocus"
diffusers_script_path = diffusers_folder / "lollms_diffusers.py" fooocus_script_path = fooocus_folder / "lollms_fooocus.py"
git_pull(diffusers_folder) git_pull(fooocus_folder)
if diffusers_script_path.exists(): if fooocus_script_path.exists():
ASCIIColors.success("lollms_diffusers found.") ASCIIColors.success("lollms_fooocus found.")
ASCIIColors.success("Loading source file...",end="") ASCIIColors.success("Loading source file...",end="")
# use importlib to load the module from the file path # use importlib to load the module from the file path
from lollms.services.diffusers.lollms_diffusers import LollmsFooocus from lollms.services.fooocus.lollms_fooocus import LollmsFooocus
ASCIIColors.success("ok") ASCIIColors.success("ok")
return LollmsFooocus return LollmsFooocus
@ -169,7 +163,8 @@ class LollmsFooocus(LollmsTTI):
): ):
if output_path is None: if output_path is None:
output_path = self.output_dir output_path = self.output_dir
from diffusers.utils.pil_utils import pt_to_pil
self.client.predict()
image = self.model(positive_prompt, negative_prompt=negative_prompt, guidance_scale=scale, num_inference_steps=steps,).images[0] image = self.model(positive_prompt, negative_prompt=negative_prompt, guidance_scale=scale, num_inference_steps=steps,).images[0]
output_path = Path(output_path) output_path = Path(output_path)
fn = find_next_available_filename(output_path,"diff_img_") fn = find_next_available_filename(output_path,"diff_img_")