import os import gc import gradio as gr import numpy as np import torch import json import spaces import config import utils import logging from PIL import Image, PngImagePlugin from datetime import datetime from diffusers.models import AutoencoderKL from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) DESCRIPTION = "Animagine XL 3.1" if not torch.cuda.is_available(): DESCRIPTION += "\n

Running on CPU 🥶 This demo does not work on CPU.

" IS_COLAB = utils.is_google_colab() or os.getenv("IS_COLAB") == "1" HF_TOKEN = os.getenv("HF_TOKEN") CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1" MIN_IMAGE_SIZE = int(os.getenv("MIN_IMAGE_SIZE", "512")) MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048")) USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1" ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1" OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./outputs") MODEL = os.getenv( "MODEL", "https://huggingface.co/cagliostrolab/animagine-xl-3.1/blob/main/animagine-xl-3.1.safetensors", ) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") def load_pipeline(model_name): vae = AutoencoderKL.from_pretrained( "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, ) pipeline = ( StableDiffusionXLPipeline.from_single_file if MODEL.endswith(".safetensors") else StableDiffusionXLPipeline.from_pretrained ) img_pipeline = ( StableDiffusionXLImg2ImgPipeline.from_single_file if MODEL.endswith(".safetensors") else StableDiffusionXLImg2ImgPipeline.from_pretrained ) pipe = pipeline( model_name, vae=vae, torch_dtype=torch.float16, custom_pipeline="lpw_stable_diffusion_xl", use_safetensors=True, add_watermarker=False, use_auth_token=HF_TOKEN, ) img_pipe = img_pipeline( model_name, vae=vae, torch_dtype=torch.float16, custom_pipeline="lpw_stable_diffusion_xl", use_safetensors=True, add_watermarker=False, use_auth_token=HF_TOKEN, ) pipe.to(device) img_pipe.to(device) return pipe, img_pipe def load_img(resize_width,img: str): img = Image.open(img) width, height = img.size scale = resize_width / width resize_height = int(height * scale) img = img.resize((resize_width, resize_height), Image.Resampling.LANCZOS) return img, resize_width, resize_height @spaces.GPU def generate( prompt: str, negative_prompt: str = "", seed: int = 0, custom_width: int = 1024, custom_height: int = 1024, guidance_scale: float = 7.0, num_inference_steps: int = 28, sampler: str = "Euler a", aspect_ratio_selector: str = "896 x 1152", style_selector: str = "(None)", quality_selector: str = "Standard v3.1", use_upscaler: bool = False, upscaler_strength: float = 0.55, upscale_by: float = 1.5, add_quality_tags: bool = True, isImg2Img: bool = True, img_path: str= "", img2img_strength: float=0.65, progress=gr.Progress(track_tqdm=True), ): generator = utils.seed_everything(seed) width, height = utils.aspect_ratio_handler( aspect_ratio_selector, custom_width, custom_height, ) prompt = utils.add_wildcard(prompt, wildcard_files) prompt, negative_prompt = utils.preprocess_prompt( quality_prompt, quality_selector, prompt, negative_prompt, add_quality_tags ) prompt, negative_prompt = utils.preprocess_prompt( styles, style_selector, prompt, negative_prompt ) width, height = utils.preprocess_image_dimensions(width, height) backup_scheduler = pipe.scheduler pipe.scheduler = utils.get_scheduler(pipe.scheduler.config, sampler) img_backup_scheduler = img_pipe.scheduler img_pipe.scheduler = utils.get_scheduler(img_pipe.scheduler.config, sampler) if use_upscaler: upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components) metadata = { "prompt": prompt, "negative_prompt": negative_prompt, "resolution": f"{width} x {height}", "guidance_scale": guidance_scale, "num_inference_steps": num_inference_steps, "seed": seed, "sampler": sampler, "sdxl_style": style_selector, "add_quality_tags": add_quality_tags, "quality_tags": quality_selector, "isImg2Img": isImg2Img, "img_path": img_path, "img2img_strength": img2img_strength } if use_upscaler: new_width = int(width * upscale_by) new_height = int(height * upscale_by) metadata["use_upscaler"] = { "upscale_method": "nearest-exact", "upscaler_strength": upscaler_strength, "upscale_by": upscale_by, "new_resolution": f"{new_width} x {new_height}", } else: metadata["use_upscaler"] = None metadata["Model"] = { "Model": DESCRIPTION, "Model hash": "e3c47aedb0", } logger.info(json.dumps(metadata, indent=4)) try: if use_upscaler: if isImg2Img: print("Img2Img") img, img_width, img_height = load_img(1024, img_path) latents = img_pipe( prompt=prompt, negative_prompt=negative_prompt, width=img_width, height=img_height, image=img, strength=img2img_strength, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, generator=generator, output_type="latent", ).images upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by) images = upscaler_pipe( prompt=prompt, negative_prompt=negative_prompt, image=upscaled_latents, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, strength=upscaler_strength, generator=generator, output_type="pil", ).images else: latents = pipe( prompt=prompt, negative_prompt=negative_prompt, width=width, height=height, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, generator=generator, output_type="latent", ).images upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by) images = upscaler_pipe( prompt=prompt, negative_prompt=negative_prompt, image=upscaled_latents, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, strength=upscaler_strength, generator=generator, output_type="pil", ).images else: if isImg2Img: print("Img2Img") img, img_width, img_height = load_img(512, img_path) images = img_pipe( prompt=prompt, negative_prompt=negative_prompt, width=img_width, height=img_height, image=img, strength=img2img_strength, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, generator=generator, output_type="pil", ).images else: images = pipe( prompt=prompt, negative_prompt=negative_prompt, width=width, height=height, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, generator=generator, output_type="pil", ).images if images: image_paths = [ utils.save_image(image, metadata, OUTPUT_DIR, IS_COLAB) for image in images ] for image_path in image_paths: logger.info(f"Image saved as {image_path} with metadata") return image_paths, metadata except Exception as e: logger.exception(f"An error occurred: {e}") raise finally: if use_upscaler: del upscaler_pipe if isImg2Img: img_pipe.scheduler = img_backup_scheduler else: pipe.scheduler = backup_scheduler utils.free_memory() def fake_generate(*args,use_upscaler=False,**kwargs): args = ",".join(args) #result, metadata = generate(args, use_upscaler=use_upscaler) return None, None if torch.cuda.is_available(): pipe, img_pipe = load_pipeline(MODEL) logger.info("Loaded on Device!") else: pipe, img_pipe = None, None styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in config.style_list} quality_prompt = { k["name"]: (k["prompt"], k["negative_prompt"]) for k in config.quality_prompt_list } wildcard_files = utils.load_wildcard_files("wildcard") with gr.Blocks(css="style.css", theme="NoCrypt/miku@1.2.1") as demo: title = gr.HTML( f"""

{DESCRIPTION}

""", elem_id="title", ) gr.Markdown( f"""Gradio demo for [cagliostrolab/animagine-xl-3.1](https://huggingface.co/cagliostrolab/animagine-xl-3.1)""", elem_id="subtitle", ) gr.DuplicateButton( value="Duplicate Space for private use", elem_id="duplicate-button", visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1", ) with gr.Row(): with gr.Column(scale=2): with gr.Tab("Input"): with gr.Group(): prompt = gr.Text( label="Prompt", max_lines=5, placeholder="Enter your prompt", ) negative_prompt = gr.Text( label="Negative Prompt", max_lines=5, placeholder="Enter a negative prompt", ) with gr.Accordion(label="Img2Img", open=False): isImg2Img = gr.Checkbox( label="Enable Img2Img", value=False ) image = gr.Image( sources=["upload", "webcam", "clipboard"], type="filepath", #visible=False, label="Choose Image" ) img2img_strength = gr.Slider( minimum=0.05, maximum=1, step=0.05, value=0.65, label="Strength", #visible=False ) with gr.Accordion(label="Quality Tags", open=True): add_quality_tags = gr.Checkbox( label="Add Quality Tags", value=True ) quality_selector = gr.Dropdown( label="Quality Tags Presets", interactive=True, choices=list(quality_prompt.keys()), value="Standard v3.1", ) with gr.Tab("ControlNet"): pass with gr.Tab("Advanced Settings"): with gr.Group(): style_selector = gr.Radio( label="Style Preset", container=True, interactive=True, choices=list(styles.keys()), value="(None)", ) with gr.Group(): aspect_ratio_selector = gr.Radio( label="Aspect Ratio", choices=config.aspect_ratios, value="896 x 1152", container=True, ) with gr.Group(visible=False) as custom_resolution: with gr.Row(): custom_width = gr.Slider( label="Width", minimum=MIN_IMAGE_SIZE, maximum=MAX_IMAGE_SIZE, step=8, value=1024, ) custom_height = gr.Slider( label="Height", minimum=MIN_IMAGE_SIZE, maximum=MAX_IMAGE_SIZE, step=8, value=1024, ) with gr.Group(): use_upscaler = gr.Checkbox(label="Use Upscaler", value=False) with gr.Row() as upscaler_row: upscaler_strength = gr.Slider( label="Strength", minimum=0, maximum=1, step=0.05, value=0.55, visible=False, ) upscale_by = gr.Slider( label="Upscale by", minimum=1, maximum=1.5, step=0.1, value=1.5, visible=False, ) with gr.Group(): sampler = gr.Dropdown( label="Sampler", choices=config.sampler_list, interactive=True, value="Euler a", ) with gr.Group(): seed = gr.Slider( label="Seed", minimum=0, maximum=utils.MAX_SEED, step=1, value=0 ) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) with gr.Group(): with gr.Row(): guidance_scale = gr.Slider( label="Guidance scale", minimum=1, maximum=12, step=0.1, value=7.0, ) num_inference_steps = gr.Slider( label="Number of inference steps", minimum=1, maximum=50, step=1, value=28, ) with gr.Column(scale=3): with gr.Blocks(): run_button = gr.Button("Generate", variant="primary") result = gr.Gallery( label="Result", columns=1, height='100%', preview=True, show_label=False ) with gr.Accordion(label="Generation Parameters", open=False): gr_metadata = gr.JSON(label="metadata", show_label=False) gr.Examples( examples=config.examples, inputs=prompt, outputs=[result, gr_metadata], fn=lambda *args, **kwargs: fake_generate(*args, use_upscaler=True, **kwargs), cache_examples=CACHE_EXAMPLES, ) use_upscaler.change( fn=lambda x: [gr.update(visible=x), gr.update(visible=x)], inputs=use_upscaler, outputs=[upscaler_strength, upscale_by], queue=False, api_name=False, ) #isImg2Img.change( # fn=lambda x: [gr.update(visible=x), gr.update(visible=x)], # inputs=isImg2Img, # outputs=[image, img2img_strength], # queue=False, # api_name=False, #) aspect_ratio_selector.change( fn=lambda x: gr.update(visible=x == "Custom"), inputs=aspect_ratio_selector, outputs=custom_resolution, queue=False, api_name=False, ) gr.on( triggers=[ prompt.submit, negative_prompt.submit, run_button.click, ], fn=utils.randomize_seed_fn, inputs=[seed, randomize_seed], outputs=seed, queue=False, api_name=False, ).then( fn=generate, inputs=[ prompt, negative_prompt, seed, custom_width, custom_height, guidance_scale, num_inference_steps, sampler, aspect_ratio_selector, style_selector, quality_selector, use_upscaler, upscaler_strength, upscale_by, add_quality_tags, isImg2Img, image, img2img_strength ], outputs=[result, gr_metadata], api_name="run", ) if __name__ == "__main__": demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)