Update demo_gradio.py

demo_gradio.py

@@ -1,3 +1,4 @@
+
 import spaces
 import huggingface_hub
 
@@ -26,7 +27,7 @@ import glob
 import torch
 import cv2
 import argparse
-
+from diffusers.models.attention_processor import AttnProcessor2_0
 import DPT.util.io
 
 from torchvision.transforms import Compose
@@ -38,7 +39,7 @@ from DPT.dpt.transforms import Resize, NormalizeImage, PrepareForNet
 """
 Get ZeST Ready
 """
-base_model_path = "stabilityai/stable-diffusion-xl-base-1.0"
+base_model_path = "Lykon/dreamshaper-xl-lightning"
 image_encoder_path = "models/image_encoder"
 ip_ckpt = "sdxl_models/ip-adapter_sdxl_vit-h.bin"
 controlnet_path = "diffusers/controlnet-depth-sdxl-1.0"
@@ -55,7 +56,7 @@ pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(
     add_watermarker=False,
 ).to(device)
 pipe.unet = register_cross_attention_hook(pipe.unet)
-
+pipe.unet.set_attn_processor(AttnProcessor2_0())
 ip_model = IPAdapterXL(pipe, image_encoder_path, ip_ckpt, device)
 
 
@@ -161,7 +162,7 @@ def greet(input_image, material_exemplar):
     
     
     num_samples = 1
-    images = ip_model.generate(pil_image=ip_image, image=init_img, control_image=depth_map, mask_image=mask, controlnet_conditioning_scale=0.9, num_samples=num_samples, num_inference_steps=30, seed=42)
+    images = ip_model.generate(guidance_scale=2, pil_image=ip_image, image=init_img, control_image=depth_map, mask_image=mask, controlnet_conditioning_scale=0.9, num_samples=num_samples, num_inference_steps=4, seed=42)
     
     return images[0]
 
@@ -193,4 +194,4 @@ with gr.Blocks(css=css) as demo:
                 output_image = gr.Image(label="transfer result")
     submit_btn.click(fn=greet, inputs=[input_image, input_image2], outputs=[output_image])
 
-demo.queue().launch()
+demo.queue().launch()