renvx0/PS1-Graphics

1import torch2from diffusers import StableDiffusionXLPipeline3import numpy as np4import gradio as gr5import random6from compel import Compel, ReturnedEmbeddingsType7 8device = "cuda" if torch.cuda.is_available() else "cpu"9 10if torch.cuda.is_available():11  torch.cuda.max_memory_allocated(device=device)12  pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)13  pipe = pipe.to(device)14else:15  pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)16  pipe = pipe.to(device)17 18pipe.safety_checker = None19 20pipe.load_lora_weights("artificialguybr/ps1redmond-ps1-game-graphics-lora-for-sdxl", weight_name="PS1Redmond-PS1Game-Playstation1Graphics.safetensors")21lora_activation_words = "playstation 1 graphics, PS1 Game, "22 23MAX_SEED = np.iinfo(np.int32).max24MAX_IMAGE_SIZE = 102425 26def infer(conditioning, pooled, neg_conditioning, neg_pooled, height, width, num_inference_steps, guidance_scale, seed, randomize_seed, lora_weight):27  if randomize_seed:28    seed = random.randint(0, MAX_SEED)29 30  generator = torch.Generator().manual_seed(seed)31 32  image = pipe(33    prompt_embeds=conditioning, 34    pooled_prompt_embeds=pooled,35    negative_prompt_embeds=neg_conditioning, 36    negative_pooled_prompt_embeds=neg_pooled,37    height=height,38    width=width,39    num_inference_steps=num_inference_steps,40    guidance_scale=guidance_scale,41    generator=generator,42    cross_attention_kwargs={"scale": lora_weight}43  ).images[0]44 45  return image46 47def get_embeds(prompt, negative_prompt, height, width, num_inference_steps, guidance_scale, seed, randomize_seed, lora_weight):48 49  compel = Compel(50    tokenizer=[pipe.tokenizer, pipe.tokenizer_2] ,51    text_encoder=[pipe.text_encoder, pipe.text_encoder_2],52    returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,53    requires_pooled=[False, True]54  )55 56  prompt = lora_activation_words + prompt57 58  conditioning, pooled = compel(prompt)59  neg_conditioning, neg_pooled = compel(negative_prompt)60 61  image = infer(conditioning, pooled, neg_conditioning, neg_pooled, height, width, num_inference_steps, guidance_scale, seed, randomize_seed, lora_weight)62 63  return image64 65css="""66#col-container {67    margin: 0 auto;68    max-width: 520px;69}70"""71 72with gr.Blocks(css=css) as demo:73    74    with gr.Column(elem_id="col-container"):75        gr.Markdown(f"""76        # Text-to-Image Gradio Template77        Currently running on {device.upper()}.78        """)79        80        with gr.Row():81            82            prompt = gr.Text(83                label="Prompt",84                show_label=False,85                max_lines=1,86                placeholder="Enter your prompt",87                container=False,88            )89            90            run_button = gr.Button("Run", scale=0)91        92        result = gr.Image(label="Result", show_label=False)93 94        with gr.Accordion("Advanced Settings", open=False):95            96            negative_prompt = gr.Text(97                label="Negative prompt",98                max_lines=1,99                placeholder="Enter a negative prompt",100                visible=True,101            )102            103            seed = gr.Slider(104                label="Seed",105                minimum=0,106                maximum=MAX_SEED,107                step=1,108                value=0,109            )110            111            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)112            113            with gr.Row():114                115                width = gr.Slider(116                    label="Width",117                    minimum=256,118                    maximum=MAX_IMAGE_SIZE,119                    step=32,120                    value=1024,121                )122                123                height = gr.Slider(124                    label="Height",125                    minimum=256,126                    maximum=MAX_IMAGE_SIZE,127                    step=32,128                    value=1024,129                )130            131            with gr.Row():132                133                guidance_scale = gr.Slider(134                    label="Guidance scale",135                    minimum=0.0,136                    maximum=10.0,137                    step=0.1,138                    value=7.5,139                )140                141                num_inference_steps = gr.Slider(142                    label="Number of inference steps",143                    minimum=1,144                    maximum=100,145                    step=1,146                    value=30,147                )148              149            with gr.Row():150 151                lora_weight = gr.Slider(152                    label="LoRA weight",153                    minimum=0.0,154                    maximum=5.0,155                    step=0.01,156                    value=1,157                )158 159    run_button.click(160        fn = get_embeds,161        inputs = [prompt, negative_prompt, height, width, num_inference_steps, guidance_scale, seed, randomize_seed, lora_weight],162        outputs = [result]163    )164 165demo.launch(debug=True)