Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python | |
| from __future__ import annotations | |
| import os | |
| import random | |
| import toml | |
| import gradio as gr | |
| import numpy as np | |
| import PIL.Image | |
| import torch | |
| import utils | |
| import gc | |
| from safetensors.torch import load_file | |
| import lora_diffusers | |
| from lora_diffusers import LoRANetwork, create_network_from_weights | |
| from huggingface_hub import hf_hub_download | |
| from diffusers.models import AutoencoderKL | |
| from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler | |
| DESCRIPTION = "Animagine XL" | |
| if not torch.cuda.is_available(): | |
| DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>" | |
| IS_COLAB = utils.is_google_colab() | |
| MAX_SEED = np.iinfo(np.int32).max | |
| CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1" | |
| 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" | |
| MODEL = "Linaqruf/animagine-xl" | |
| device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
| if torch.cuda.is_available(): | |
| pipe = DiffusionPipeline.from_pretrained( | |
| MODEL, | |
| torch_dtype=torch.float16, | |
| custom_pipeline="lpw_stable_diffusion_xl.py", | |
| use_safetensors=True, | |
| variant="fp16", | |
| ) | |
| pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config) | |
| if ENABLE_CPU_OFFLOAD: | |
| pipe.enable_model_cpu_offload() | |
| else: | |
| pipe.to(device) | |
| if USE_TORCH_COMPILE: | |
| pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) | |
| else: | |
| pipe = None | |
| def randomize_seed_fn(seed: int, randomize_seed: bool) -> int: | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) | |
| return seed | |
| def get_image_path(base_path): | |
| extensions = [".jpg", ".jpeg", ".png", ".bmp", ".gif"] | |
| for ext in extensions: | |
| if os.path.exists(base_path + ext): | |
| return base_path + ext | |
| # If no match is found, return None or raise an error | |
| return None | |
| def update_selection(selected_state: gr.SelectData): | |
| lora_repo = sdxl_loras[selected_state.index]["repo"] | |
| lora_weight = sdxl_loras[selected_state.index]["multiplier"] | |
| updated_selected_info = f"{lora_repo}" | |
| updated_prompt = sdxl_loras[selected_state.index]["sample_prompt"] | |
| updated_negative = sdxl_loras[selected_state.index]["sample_negative"] | |
| return ( | |
| updated_selected_info, | |
| selected_state, | |
| lora_weight, | |
| updated_prompt, | |
| negative_presets_dict.get(updated_negative, ""), | |
| updated_negative, | |
| ) | |
| def create_network(text_encoders, unet, state_dict, multiplier, device): | |
| network = create_network_from_weights( | |
| text_encoders, unet, state_dict, multiplier=multiplier | |
| ) | |
| network.load_state_dict(state_dict) | |
| network.to(device, dtype=unet.dtype) | |
| network.apply_to(multiplier=multiplier) | |
| return network | |
| # def backup_sd(state_dict): | |
| # for k, v in state_dict.items(): | |
| # state_dict[k] = v.detach().cpu() | |
| # return state_dict | |
| def generate( | |
| prompt: str, | |
| negative_prompt: str = "", | |
| prompt_2: str = "", | |
| negative_prompt_2: str = "", | |
| use_prompt_2: bool = False, | |
| seed: int = 0, | |
| width: int = 1024, | |
| height: int = 1024, | |
| target_width: int = 1024, | |
| target_height: int = 1024, | |
| original_width: int = 4096, | |
| original_height: int = 4096, | |
| guidance_scale: float = 12.0, | |
| num_inference_steps: int = 50, | |
| use_lora: bool = False, | |
| lora_weight: float = 1.0, | |
| set_target_size: bool = False, | |
| set_original_size: bool = False, | |
| selected_state: str = "", | |
| ) -> PIL.Image.Image: | |
| generator = torch.Generator().manual_seed(seed) | |
| network = None # Initialize to None | |
| network_state = {"current_lora": None, "multiplier": None} | |
| # _unet = pipe.unet.state_dict() | |
| # backup_sd(_unet) | |
| # _text_encoder = pipe.text_encoder.state_dict() | |
| # backup_sd(_text_encoder) | |
| # _text_encoder_2 = pipe.text_encoder_2.state_dict() | |
| # backup_sd(_text_encoder_2) | |
| if not set_original_size: | |
| original_width = 4096 | |
| original_height = 4096 | |
| if not set_target_size: | |
| target_width = width | |
| target_height = height | |
| if negative_prompt == "": | |
| negative_prompt = None | |
| if not use_prompt_2: | |
| prompt_2 = None | |
| negative_prompt_2 = None | |
| if negative_prompt_2 == "": | |
| negative_prompt_2 = None | |
| if use_lora: | |
| if not selected_state: | |
| raise Exception("You must select a LoRA") | |
| repo_name = sdxl_loras[selected_state.index]["repo"] | |
| full_path_lora = saved_names[selected_state.index] | |
| weight_name = sdxl_loras[selected_state.index]["weights"] | |
| lora_sd = load_file(full_path_lora) | |
| text_encoders = [pipe.text_encoder, pipe.text_encoder_2] | |
| if network_state["current_lora"] != repo_name: | |
| network = create_network( | |
| text_encoders, pipe.unet, lora_sd, lora_weight, device | |
| ) | |
| network_state["current_lora"] = repo_name | |
| network_state["multiplier"] = lora_weight | |
| elif network_state["multiplier"] != lora_weight: | |
| network = create_network( | |
| text_encoders, pipe.unet, lora_sd, lora_weight, device | |
| ) | |
| network_state["multiplier"] = lora_weight | |
| else: | |
| if network: | |
| network.unapply_to() | |
| network = None | |
| network_state = {"current_lora": None, "multiplier": None} | |
| try: | |
| image = pipe( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| prompt_2=prompt_2, | |
| negative_prompt_2=negative_prompt_2, | |
| width=width, | |
| height=height, | |
| target_size=(target_width, target_height), | |
| original_size=(original_width, original_height), | |
| guidance_scale=guidance_scale, | |
| num_inference_steps=num_inference_steps, | |
| generator=generator, | |
| output_type="pil", | |
| ).images[0] | |
| if network: | |
| network.unapply_to() | |
| network = None | |
| return image | |
| except Exception as e: | |
| print(f"An error occurred: {e}") | |
| raise | |
| finally: | |
| # pipe.unet.load_state_dict(_unet) | |
| # pipe.text_encoder.load_state_dict(_text_encoder) | |
| # pipe.text_encoder_2.load_state_dict(_text_encoder_2) | |
| # del _unet, _text_encoder, _text_encoder_2 | |
| if network: | |
| network.unapply_to() | |
| network = None | |
| if use_lora: | |
| del lora_sd, text_encoders | |
| gc.collect() | |
| examples = [ | |
| "face focus, cute, masterpiece, best quality, 1girl, green hair, sweater, looking at viewer, upper body, beanie, outdoors, night, turtleneck", | |
| "face focus, bishounen, masterpiece, best quality, 1boy, green hair, sweater, looking at viewer, upper body, beanie, outdoors, night, turtleneck", | |
| ] | |
| negative_presets_dict = { | |
| "None": "", | |
| "Standard": "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry", | |
| "Weighted": "(low quality, worst quality:1.2), 3d, watermark, signature, ugly, poorly drawn, bad image", | |
| } | |
| with open("lora.toml", "r") as file: | |
| data = toml.load(file) | |
| sdxl_loras = [ | |
| { | |
| "image": get_image_path(item["image"]), | |
| "title": item["title"], | |
| "repo": item["repo"], | |
| "weights": item["weights"], | |
| "multiplier": item["multiplier"] if "multiplier" in item else "1.0", | |
| "sample_prompt": item["sample_prompt"], | |
| "sample_negative": item["sample_negative"], | |
| } | |
| for item in data["data"] | |
| ] | |
| saved_names = [hf_hub_download(item["repo"], item["weights"]) for item in sdxl_loras] | |
| with gr.Blocks(css="style.css", theme="NoCrypt/miku@1.2.1") as demo: | |
| title = gr.HTML( | |
| f"""<h1><span>{DESCRIPTION}</span></h1>""", | |
| elem_id="title", | |
| ) | |
| gr.Markdown( | |
| f"""Gradio demo for [Linaqruf/animagine-xl](https://huggingface.co/Linaqruf/Animagine-XL)""", | |
| elem_id="subtitle", | |
| ) | |
| gr.DuplicateButton( | |
| value="Duplicate Space for private use", | |
| elem_id="duplicate-button", | |
| visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1", | |
| ) | |
| selected_state = gr.State() | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| 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", | |
| value="lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry", | |
| ) | |
| with gr.Accordion(label="Negative Presets", open=False): | |
| negative_presets = gr.Dropdown( | |
| label="Negative Presets", | |
| show_label=False, | |
| choices=list(negative_presets_dict.keys()), | |
| value="Standard", | |
| ) | |
| with gr.Row(): | |
| use_prompt_2 = gr.Checkbox(label="Use prompt 2", value=False) | |
| use_lora = gr.Checkbox(label="Use LoRA", value=False) | |
| with gr.Group(visible=False) as prompt2_group: | |
| prompt_2 = gr.Text( | |
| label="Prompt 2", | |
| max_lines=5, | |
| placeholder="Enter your prompt", | |
| ) | |
| negative_prompt_2 = gr.Text( | |
| label="Negative prompt 2", | |
| max_lines=5, | |
| placeholder="Enter a negative prompt", | |
| ) | |
| with gr.Group(visible=False) as lora_group: | |
| selector_info = gr.Text( | |
| label="Selected LoRA", | |
| max_lines=1, | |
| value="No LoRA selected.", | |
| ) | |
| lora_selection = gr.Gallery( | |
| value=[(item["image"], item["title"]) for item in sdxl_loras], | |
| label="Animagine XL LoRA", | |
| show_label=False, | |
| allow_preview=False, | |
| columns=2, | |
| elem_id="gallery", | |
| show_share_button=False, | |
| ) | |
| lora_weight = gr.Slider( | |
| label="Multiplier", | |
| minimum=0, | |
| maximum=1, | |
| step=0.05, | |
| value=1, | |
| ) | |
| with gr.Group(): | |
| with gr.Row(): | |
| width = gr.Slider( | |
| label="Width", | |
| minimum=256, | |
| maximum=MAX_IMAGE_SIZE, | |
| step=32, | |
| value=1024, | |
| ) | |
| height = gr.Slider( | |
| label="Height", | |
| minimum=256, | |
| maximum=MAX_IMAGE_SIZE, | |
| step=32, | |
| value=1024, | |
| ) | |
| with gr.Accordion(label="Advanced Options", open=False): | |
| seed = gr.Slider( | |
| label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0 | |
| ) | |
| randomize_seed = gr.Checkbox(label="Randomize seed", value=True) | |
| with gr.Row(): | |
| guidance_scale = gr.Slider( | |
| label="Guidance scale", | |
| minimum=1, | |
| maximum=20, | |
| step=0.1, | |
| value=12.0, | |
| ) | |
| num_inference_steps = gr.Slider( | |
| label="Number of inference steps", | |
| minimum=10, | |
| maximum=100, | |
| step=1, | |
| value=50, | |
| ) | |
| with gr.Group(): | |
| with gr.Row(): | |
| set_target_size = gr.Checkbox( | |
| label="Target Size", value=False | |
| ) | |
| set_original_size = gr.Checkbox( | |
| label="Original Size", value=False | |
| ) | |
| with gr.Group(): | |
| with gr.Row(): | |
| original_width = gr.Slider( | |
| label="Original Width", | |
| minimum=1024, | |
| maximum=4096, | |
| step=32, | |
| value=4096, | |
| visible=False, | |
| ) | |
| original_height = gr.Slider( | |
| label="Original Height", | |
| minimum=1024, | |
| maximum=4096, | |
| step=32, | |
| value=4096, | |
| visible=False, | |
| ) | |
| with gr.Row(): | |
| target_width = gr.Slider( | |
| label="Target Width", | |
| minimum=1024, | |
| maximum=4096, | |
| step=32, | |
| value=width.value, | |
| visible=False, | |
| ) | |
| target_height = gr.Slider( | |
| label="Target Height", | |
| minimum=1024, | |
| maximum=4096, | |
| step=32, | |
| value=height.value, | |
| visible=False, | |
| ) | |
| with gr.Column(scale=2): | |
| with gr.Blocks(): | |
| run_button = gr.Button("Generate", variant="primary") | |
| result = gr.Image(label="Result", show_label=False) | |
| gr.Examples( | |
| examples=examples, | |
| inputs=prompt, | |
| outputs=result, | |
| fn=generate, | |
| cache_examples=CACHE_EXAMPLES, | |
| ) | |
| lora_selection.select( | |
| update_selection, | |
| outputs=[ | |
| selector_info, | |
| selected_state, | |
| lora_weight, | |
| prompt, | |
| negative_prompt, | |
| negative_presets, | |
| ], | |
| queue=False, | |
| show_progress=False, | |
| ) | |
| use_prompt_2.change( | |
| fn=lambda x: gr.update(visible=x), | |
| inputs=use_prompt_2, | |
| outputs=prompt2_group, | |
| queue=False, | |
| api_name=False, | |
| ) | |
| negative_presets.change( | |
| fn=lambda x: gr.update(value=negative_presets_dict.get(x, "")), | |
| inputs=negative_presets, | |
| outputs=negative_prompt, | |
| queue=False, | |
| api_name=False, | |
| ) | |
| use_lora.change( | |
| fn=lambda x: gr.update(visible=x), | |
| inputs=use_lora, | |
| outputs=lora_group, | |
| queue=False, | |
| api_name=False, | |
| ) | |
| set_target_size.change( | |
| fn=lambda x: (gr.update(visible=x), gr.update(visible=x)), | |
| inputs=set_target_size, | |
| outputs=[target_width, target_height], | |
| queue=False, | |
| api_name=False, | |
| ) | |
| set_original_size.change( | |
| fn=lambda x: (gr.update(visible=x), gr.update(visible=x)), | |
| inputs=set_original_size, | |
| outputs=[original_width, original_height], | |
| queue=False, | |
| api_name=False, | |
| ) | |
| width.change( | |
| fn=lambda x: gr.update(value=x), | |
| inputs=width, | |
| outputs=target_width, | |
| queue=False, | |
| api_name=False, | |
| ) | |
| height.change( | |
| fn=lambda x: gr.update(value=x), | |
| inputs=height, | |
| outputs=target_height, | |
| queue=False, | |
| api_name=False, | |
| ) | |
| inputs = [ | |
| prompt, | |
| negative_prompt, | |
| prompt_2, | |
| negative_prompt_2, | |
| use_prompt_2, | |
| seed, | |
| width, | |
| height, | |
| target_width, | |
| target_height, | |
| original_width, | |
| original_height, | |
| guidance_scale, | |
| num_inference_steps, | |
| use_lora, | |
| lora_weight, | |
| set_target_size, | |
| set_original_size, | |
| selected_state, | |
| ] | |
| prompt.submit( | |
| fn=randomize_seed_fn, | |
| inputs=[seed, randomize_seed], | |
| outputs=seed, | |
| queue=False, | |
| api_name=False, | |
| ).then( | |
| fn=generate, | |
| inputs=inputs, | |
| outputs=result, | |
| api_name="run", | |
| ) | |
| negative_prompt.submit( | |
| fn=randomize_seed_fn, | |
| inputs=[seed, randomize_seed], | |
| outputs=seed, | |
| queue=False, | |
| api_name=False, | |
| ).then( | |
| fn=generate, | |
| inputs=inputs, | |
| outputs=result, | |
| api_name=False, | |
| ) | |
| prompt_2.submit( | |
| fn=randomize_seed_fn, | |
| inputs=[seed, randomize_seed], | |
| outputs=seed, | |
| queue=False, | |
| api_name=False, | |
| ).then( | |
| fn=generate, | |
| inputs=inputs, | |
| outputs=result, | |
| api_name=False, | |
| ) | |
| negative_prompt_2.submit( | |
| fn=randomize_seed_fn, | |
| inputs=[seed, randomize_seed], | |
| outputs=seed, | |
| queue=False, | |
| api_name=False, | |
| ).then( | |
| fn=generate, | |
| inputs=inputs, | |
| outputs=result, | |
| api_name=False, | |
| ) | |
| run_button.click( | |
| fn=randomize_seed_fn, | |
| inputs=[seed, randomize_seed], | |
| outputs=seed, | |
| queue=False, | |
| api_name=False, | |
| ).then( | |
| fn=generate, | |
| inputs=inputs, | |
| outputs=result, | |
| api_name=False, | |
| ) | |
| demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB) | |