Update app.py

app.py

@@ -1,45 +1,66 @@
-import gradio as gr
+import torch
+from diffusers import StableDiffusionXLPipeline
 import numpy as np
+import gradio as gr
 import random
-from diffusers import AutoPipelineForImage2Image
-from diffusers.utils import make_image_grid, load_image
+from compel import Compel, ReturnedEmbeddingsType
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 if torch.cuda.is_available():
-    torch.cuda.max_memory_allocated(device=device)
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to(device)
-else: 
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", use_safetensors=True)
-    pipe = pipe.to(device)
+  torch.cuda.max_memory_allocated(device=device)
+  pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
+  pipe = pipe.to(device)
+else:
+  pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
+  pipe = pipe.to(device)
+
+pipe.safety_checker = None
+
+pipe.load_lora_weights("artificialguybr/ps1redmond-ps1-game-graphics-lora-for-sdxl", weight_name="PS1Redmond-PS1Game-Playstation1Graphics.safetensors")
+lora_activation_words = "playstation 1 graphics, PS1 Game, "
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-pipe.load_lora_weights("artificialguybr/ps1redmond-ps1-game-graphics-lora-for-sdxl")
+def infer(conditioning, pooled, neg_conditioning, neg_pooled, height, width, num_inference_steps, guidance_scale, seed, randomize_seed, lora_weight):
+  if randomize_seed:
+    seed = random.randint(0, MAX_SEED)
 
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+  generator = torch.Generator().manual_seed(seed)
 
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-        
-    generator = torch.Generator().manual_seed(seed)
+  image = pipe(
+    prompt_embeds=conditioning, 
+    pooled_prompt_embeds=pooled,
+    negative_prompt_embeds=neg_conditioning, 
+    negative_pooled_prompt_embeds=neg_pooled,
+    height=height,
+    width=width,
+    num_inference_steps=num_inference_steps,
+    guidance_scale=guidance_scale,
+    generator=generator,
+    cross_attention_kwargs={"scale": lora_weight}
+  ).images[0]
 
-    url = "https://cmmodels.de/wp-content/uploads/2021/08/gabriel-new-face-white-shirt-brown-hair-pretty-boy-young-white-background-studio-brown-hair-wet-look.jpg"
-    init_image = load_image(url)
-    
-    image = pipeline(prompt, image=init_image, strength=0.5).images[0]
-    grid = make_image_grid([init_image, image], rows=1, cols=2)
-    
-    return grid
+  return image
+
+def get_embeds(prompt, negative_prompt, height, width, num_inference_steps, guidance_scale, seed, randomize_seed, lora_weight):
+
+  compel = Compel(
+    tokenizer=[pipe.tokenizer, pipe.tokenizer_2] ,
+    text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
+    returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+    requires_pooled=[False, True]
+  )
+
+  prompt = lora_activation_words + prompt
 
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
+  conditioning, pooled = compel(prompt)
+  neg_conditioning, neg_pooled = compel(negative_prompt)
+
+  image = infer(conditioning, pooled, neg_conditioning, neg_pooled, height, width, num_inference_steps, guidance_scale, seed, randomize_seed, lora_weight)
+
+  return image
 
 css="""
 #col-container {
@@ -48,17 +69,12 @@ css="""
 }
 """
 
-if torch.cuda.is_available():
-    power_device = "GPU"
-else:
-    power_device = "CPU"
-
 with gr.Blocks(css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""
         # Text-to-Image Gradio Template
-        Currently running on {power_device}.
+        Currently running on {device.upper()}.
         """)
         
         with gr.Row():
@@ -81,7 +97,7 @@ with gr.Blocks(css=css) as demo:
                 label="Negative prompt",
                 max_lines=1,
                 placeholder="Enter a negative prompt",
-                visible=False,
+                visible=True,
             )
             
             seed = gr.Slider(
@@ -101,7 +117,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,
+                    value=1024,
                 )
                 
                 height = gr.Slider(
@@ -109,7 +125,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,
+                    value=1024,
                 )
             
             with gr.Row():
@@ -119,26 +135,31 @@ with gr.Blocks(css=css) as demo:
                     minimum=0.0,
                     maximum=10.0,
                     step=0.1,
-                    value=0.0,
+                    value=7.5,
                 )
                 
                 num_inference_steps = gr.Slider(
                     label="Number of inference steps",
                     minimum=1,
-                    maximum=12,
+                    maximum=100,
                     step=1,
-                    value=2,
+                    value=30,
+                )
+              
+            with gr.Row():
+
+                lora_weight = gr.Slider(
+                    label="LoRA weight",
+                    minimum=0.0,
+                    maximum=5.0,
+                    step=0.01,
+                    value=1,
                 )
-        
-        gr.Examples(
-            examples = examples,
-            inputs = [prompt]
-        )
 
     run_button.click(
-        fn = infer,
-        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        fn = get_embeds,
+        inputs = [prompt, negative_prompt, height, width, num_inference_steps, guidance_scale, seed, randomize_seed, lora_weight],
         outputs = [result]
     )
 
-demo.queue().launch()
+demo.launch(debug=True)