justmalhar/Animagine-XL
0
1#!/usr/bin/env python2 3from __future__ import annotations4 5import os6import random7 8import gradio as gr9import numpy as np10import PIL.Image11import torch12from diffusers.models import AutoencoderKL13from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler14 15DESCRIPTION = '# Animagine XL'16if not torch.cuda.is_available():17 DESCRIPTION += '\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>'18 19MAX_SEED = np.iinfo(np.int32).max20CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv(21 'CACHE_EXAMPLES') == '1'22MAX_IMAGE_SIZE = int(os.getenv('MAX_IMAGE_SIZE', '2048'))23USE_TORCH_COMPILE = os.getenv('USE_TORCH_COMPILE') == '1'24ENABLE_CPU_OFFLOAD = os.getenv('ENABLE_CPU_OFFLOAD') == '1'25 26MODEL = "Linaqruf/animagine-xl"27 28device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')29if torch.cuda.is_available():30 pipe = StableDiffusionXLPipeline.from_pretrained(31 MODEL,32 torch_dtype=torch.float16,33 use_safetensors=True,34 variant='fp16')35 36 pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)37 38 if ENABLE_CPU_OFFLOAD:39 pipe.enable_model_cpu_offload()40 else:41 pipe.to(device)42 43 if USE_TORCH_COMPILE:44 pipe.unet = torch.compile(pipe.unet,45 mode='reduce-overhead',46 fullgraph=True)47else:48 pipe = None49 50 51def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:52 if randomize_seed:53 seed = random.randint(0, MAX_SEED)54 return seed55 56 57def generate(prompt: str,58 negative_prompt: str = '',59 prompt_2: str = '',60 negative_prompt_2: str = '',61 use_prompt_2: bool = False,62 seed: int = 0,63 width: int = 1024,64 height: int = 1024,65 target_width: int = 1024,66 target_height: int = 1024,67 original_width: int = 4096,68 original_height: int = 4096,69 guidance_scale_base: float = 12.0,70 num_inference_steps_base: int = 50) -> PIL.Image.Image:71 72 generator = torch.Generator().manual_seed(seed)73 74 if negative_prompt == '':75 negative_prompt = None # type: ignore76 if not use_prompt_2:77 prompt_2 = None # type: ignore78 negative_prompt_2 = None # type: ignore79 if negative_prompt_2 == '':80 negative_prompt_2 = None 81 82 return pipe(prompt=prompt,83 negative_prompt=negative_prompt,84 prompt_2=prompt_2,85 negative_prompt_2=negative_prompt_2,86 width=width,87 height=height,88 target_size=(target_width, target_height),89 original_size=(original_width, original_height),90 guidance_scale=guidance_scale_base,91 num_inference_steps=num_inference_steps_base,92 generator=generator,93 output_type='pil').images[0]94 95 96examples = [97 'face focus, cute, masterpiece, best quality, 1girl, green hair, sweater, looking at viewer, upper body, beanie, outdoors, night, turtleneck',98 'face focus, bishounen, masterpiece, best quality, 1boy, green hair, sweater, looking at viewer, upper body, beanie, outdoors, night, turtleneck',99]100 101# choices = [102# "Vertical (9:16)",103# "Portrait (4:5)",104# "Square (1:1)",105# "Photo (4:3)",106# "Landscape (3:2)",107# "Widescreen (16:9)",108# "Cinematic (21:9)",109# ]110 111# choice_to_size = {112# "Vertical (9:16)": (768, 1344),113# "Portrait (4:5)": (912, 1144),114# "Square (1:1)": (1024, 1024),115# "Photo (4:3)": (1184, 888),116# "Landscape (3:2)": (1256, 832),117# "Widescreen (16:9)": (1368, 768),118# "Cinematic (21:9)": (1568, 672),119# }120 121with gr.Blocks(css='style.css') as demo:122 gr.Markdown(DESCRIPTION)123 gr.DuplicateButton(value='Duplicate Space for private use',124 elem_id='duplicate-button',125 visible=os.getenv('SHOW_DUPLICATE_BUTTON') == '1')126 with gr.Row():127 with gr.Column(scale=1):128 prompt = gr.Text(129 label='Prompt',130 max_lines=1,131 placeholder='Enter your prompt',132 )133 negative_prompt = gr.Text(134 label='Negative Prompt',135 max_lines=1,136 placeholder='Enter a negative prompt',137 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',138 )139 use_prompt_2 = gr.Checkbox(140 label='Use prompt 2', 141 value=False142 ) 143 prompt_2 = gr.Text(144 label='Prompt 2',145 max_lines=1,146 placeholder='Enter your prompt',147 visible=False,148 )149 negative_prompt_2 = gr.Text(150 label='Negative prompt 2',151 max_lines=1,152 placeholder='Enter a negative prompt',153 visible=False,154 )155 156 # with gr.Row():157 # aspect_ratio = gr.Dropdown(choices=choices, label="Aspect Ratio Preset", value=choices[2])158 with gr.Row():159 width = gr.Slider(160 label='Width',161 minimum=256,162 maximum=MAX_IMAGE_SIZE,163 step=32,164 value=1024,165 )166 height = gr.Slider(167 label='Height',168 minimum=256,169 maximum=MAX_IMAGE_SIZE,170 step=32,171 value=1024,172 )173 with gr.Accordion(label='Advanced Config', open=False):174 with gr.Accordion(label='Conditioning Resolution', open=False):175 with gr.Row():176 original_width = gr.Slider(177 label='Original Width',178 minimum=1024,179 maximum=4096,180 step=32,181 value=4096,182 )183 original_height = gr.Slider(184 label='Original Height',185 minimum=1024,186 maximum=4096,187 step=32,188 value=4096,189 )190 with gr.Row():191 target_width = gr.Slider(192 label='Target Width',193 minimum=1024,194 maximum=4096,195 step=32,196 value=1024,197 )198 target_height = gr.Slider(199 label='Target Height',200 minimum=1024,201 maximum=4096,202 step=32,203 value=1024,204 )205 seed = gr.Slider(label='Seed',206 minimum=0,207 maximum=MAX_SEED,208 step=1,209 value=0)210 211 randomize_seed = gr.Checkbox(label='Randomize seed', value=True)212 with gr.Row():213 guidance_scale_base = gr.Slider(214 label='Guidance scale',215 minimum=1,216 maximum=20,217 step=0.1,218 value=12.0)219 num_inference_steps_base = gr.Slider(220 label='Number of inference steps',221 minimum=10,222 maximum=100,223 step=1,224 value=50)225 226 with gr.Column(scale=2):227 with gr.Blocks():228 run_button = gr.Button('Generate')229 result = gr.Image(label='Result', show_label=False)230 231 gr.Examples(examples=examples,232 inputs=prompt,233 outputs=result,234 fn=generate,235 cache_examples=CACHE_EXAMPLES)236 237 use_prompt_2.change(238 fn=lambda x: gr.update(visible=x),239 inputs=use_prompt_2,240 outputs=prompt_2,241 queue=False,242 api_name=False,243 )244 use_prompt_2.change(245 fn=lambda x: gr.update(visible=x),246 inputs=use_prompt_2,247 outputs=negative_prompt_2,248 queue=False,249 api_name=False,250 )251 252 inputs = [253 prompt,254 negative_prompt,255 prompt_2,256 negative_prompt_2,257 use_prompt_2,258 seed,259 width,260 height,261 target_width,262 target_height,263 original_width,264 original_height,265 guidance_scale_base,266 num_inference_steps_base,267 ]268 prompt.submit(269 fn=randomize_seed_fn,270 inputs=[seed, randomize_seed],271 outputs=seed,272 queue=False,273 api_name=False,274 ).then(275 fn=generate,276 inputs=inputs,277 outputs=result,278 api_name='run',279 )280 negative_prompt.submit(281 fn=randomize_seed_fn,282 inputs=[seed, randomize_seed],283 outputs=seed,284 queue=False,285 api_name=False,286 ).then(287 fn=generate,288 inputs=inputs,289 outputs=result,290 api_name=False,291 )292 prompt_2.submit(293 fn=randomize_seed_fn,294 inputs=[seed, randomize_seed],295 outputs=seed,296 queue=False,297 api_name=False,298 ).then(299 fn=generate,300 inputs=inputs,301 outputs=result,302 api_name=False,303 )304 negative_prompt_2.submit(305 fn=randomize_seed_fn,306 inputs=[seed, randomize_seed],307 outputs=seed,308 queue=False,309 api_name=False,310 ).then(311 fn=generate,312 inputs=inputs,313 outputs=result,314 api_name=False,315 )316 run_button.click(317 fn=randomize_seed_fn,318 inputs=[seed, randomize_seed],319 outputs=seed,320 queue=False,321 api_name=False,322 ).then(323 fn=generate,324 inputs=inputs,325 outputs=result,326 api_name=False,327 )328 329demo.queue(max_size=20).launch()330