Update raw.py

raw.py

@@ -7,7 +7,11 @@ from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
 from transformers import T5EncoderModel
 from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
 from peft import PeftModel, PeftConfig
-
+from attention_map_diffusers import (
+    attn_maps,
+    init_pipeline,
+    save_attention_maps
+)
 import gradio as gr
 huggingface_token = os.getenv("HUGGINFACE_TOKEN")
 
@@ -46,6 +50,7 @@ pipe.load_lora_weights(adapter_id3, weight_name="lora.safetensors", adapter_name
 pipe.set_adapters(["turbo", "real", "enhance"], adapter_weights=[0.9, 0.66, 0.6])
 pipe.fuse_lora(adapter_names=["turbo", "real", "enhance"], lora_scale=1.0)
 pipe.unload_lora_weights()
+pipe = init_pipeline(pipe)
 # pipe.enable_xformers_memory_efficient_attention()
 # save to the Hub
 # pipe.push_to_hub("fused-t-r")
@@ -68,6 +73,7 @@ def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_
             height=control_image.size[1],
             width=control_image.size[0]
         ).images[0]
+        
     return image
 
 # Create Gradio interface with rows and columns