Update app.py

app.py

@@ -3,17 +3,68 @@ import numpy as np
 import random
 from diffusers import DiffusionPipeline
 import torch
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+from cog_sdxl.dataset_and_utils import TokenEmbeddingsHandler
+from cog_sdxl.no_init import no_init_or_tensor
+from diffusers.models.attention_processor import LoRAAttnProcessor2_0
 
 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/sdxl-turbo", 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/sdxl-turbo", use_safetensors=True)
-    pipe = pipe.to(device)
+pipe = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    variant="fp16",
+).to(device)
+
+unet = pipe.unet
+
+lora_path = hf_hub_download(repo_id="SvenN/sdxl-emoji", filename="lora.safetensors", repo_type="model")
+embeddings_path = hf_hub_download(repo_id="SvenN/sdxl-emoji", filename="embeddings.pti", repo_type="model")
+
+tensors = load_file(lora_path)
+unet_lora_attn_procs = {}
+name_rank_map = {}
+for tk, tv in tensors.items():
+    # up is N, d
+    tensors[tk] = tv.half()
+    if tk.endswith("up.weight"):
+        proc_name = ".".join(tk.split(".")[:-3])
+        r = tv.shape[1]
+        name_rank_map[proc_name] = r
+
+for name, attn_processor in unet.attn_processors.items():
+    cross_attention_dim = (
+        None
+        if name.endswith("attn1.processor")
+        else unet.config.cross_attention_dim
+    )
+    if name.startswith("mid_block"):
+        hidden_size = unet.config.block_out_channels[-1]
+    elif name.startswith("up_blocks"):
+        block_id = int(name[len("up_blocks.")])
+        hidden_size = list(reversed(unet.config.block_out_channels))[
+            block_id
+        ]
+    elif name.startswith("down_blocks"):
+        block_id = int(name[len("down_blocks.")])
+        hidden_size = unet.config.block_out_channels[block_id]
+    with no_init_or_tensor():
+        module = LoRAAttnProcessor2_0(
+            hidden_size=hidden_size,
+            cross_attention_dim=cross_attention_dim,
+            rank=name_rank_map[name],
+        ).half()
+    unet_lora_attn_procs[name] = module.to("cuda", non_blocking=True)
+
+unet.set_attn_processor(unet_lora_attn_procs)
+unet.load_state_dict(tensors, strict=False)
+
+handler = TokenEmbeddingsHandler(
+    [pipe.text_encoder, pipe.text_encoder_2], [pipe.tokenizer, pipe.tokenizer_2]
+)
+handler.load_embeddings(embeddings_path)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
@@ -32,7 +83,8 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
         num_inference_steps = num_inference_steps, 
         width = width, 
         height = height,
-        generator = generator
+        generator = generator,
+        cross_attention_kwargs={"scale": 0.6},
     ).images[0] 
     
     return image
@@ -50,17 +102,11 @@ 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}.
+        # SDXL Emoji running on diffusers 0.25.0
         """)
         
         with gr.Row():
@@ -103,7 +149,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,
+                    value=1024,
                 )
                 
                 height = gr.Slider(
@@ -111,7 +157,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,
+                    value=1024,
                 )
             
             with gr.Row():
@@ -121,15 +167,15 @@ 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=50,
                     step=1,
-                    value=2,
+                    value=50,
                 )
         
         gr.Examples(