Update app.py

app.py

@@ -7,35 +7,31 @@ import numpy as np
 import cv2
 import gradio as gr
 from torchvision import transforms 
+import fire
 
 controlnet = ControlNetModel.from_pretrained(
     "geyongtao/HumanWild",
     torch_dtype=torch.float16
 ).to('cuda')
 
+vae = AutoencoderKL.from_pretrained(
+    madebyollin/sdxl-vae-fp16-fix, 
+    torch_dtype=torch.float16).to("cuda")
+
 pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
     controlnet=controlnet,
+    vae=vae,
     torch_dtype=torch.float16,
-    device_map='auto',
+    use_safetensors=True,
     low_cpu_mem_usage=True,
     offload_state_dict=True,
 ).to('cuda')
+pipe.controlnet.to(memory_format=torch.channels_last)
 
-pipe.scheduler = EulerAncestralDiscreteScheduler(
-    beta_start=0.00085,
-    beta_end=0.012,
-    beta_schedule="scaled_linear",
-    num_train_timesteps=1000,
-    steps_offset=1
-)
-# pipe.enable_freeu(b1=1.1, b2=1.1, s1=0.5, s2=0.7)
 # pipe.enable_xformers_memory_efficient_attention()
 pipe.force_zeros_for_empty_prompt = False
 
-# from transformers import DPTFeatureExtractor, DPTForDepthEstimation
-# depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to("cuda")
-# feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-hybrid-midas")
 
 def resize_image(image):
     image = image.convert('RGB')
@@ -83,35 +79,38 @@ def process(normal_image, prompt, negative_prompt, num_steps, controlnet_conditi
     # depth_image = get_depth_map(input_image)
     images = generate_(prompt, negative_prompt, normal_image, num_steps, controlnet_conditioning_scale, seed)
 
-    return [depth_image, images[0]]
-
+    return [normal_image, images[0]]
 
 
-block = gr.Blocks().queue()
-
-with block:
-    gr.Markdown("## BRIA 2.2 ControlNet Depth")
-    gr.HTML('''
-      <p style="margin-bottom: 10px; font-size: 94%">
-        This is a demo for ControlNet Surface Normal that using
-        <a href="https://huggingface.co/geyongtao/HumanWild" target="_blank"> HumanWild model</a> as backbone. 
-      </p>
-    ''')
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(sources=None, type="pil") # None for upload, ctrl+v and webcam
-            prompt = gr.Textbox(label="Prompt")
-            negative_prompt = gr.Textbox(label="Negative prompt", value="Logo,Watermark,Text,Ugly,Morbid,Extra fingers,Poorly drawn hands,Mutation,Blurry,Extra limbs,Gross proportions,Missing arms,Mutated hands,Long neck,Duplicate,Mutilated,Mutilated hands,Poorly drawn face,Deformed,Bad anatomy,Cloned face,Malformed limbs,Missing legs,Too many fingers")
-            num_steps = gr.Slider(label="Number of steps", minimum=25, maximum=100, value=50, step=1)
-            controlnet_conditioning_scale = gr.Slider(label="ControlNet conditioning scale", minimum=0.1, maximum=2.0, value=1.0, step=0.05)
-            seed = gr.Slider(label="Seed", minimum=0, maximum=2147483647, step=1, randomize=True,)
-            run_button = gr.Button(value="Run")
-            
-            
-        with gr.Column():
-            result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery", columns=[2], height='auto')
-    ips = [input_image, prompt, negative_prompt, num_steps, controlnet_conditioning_scale, seed]
+def run_demo():
+    block = gr.Blocks().queue()
+    
+    with block:
+        gr.Markdown("## Surface Normal ControlNet ")
+        gr.HTML('''
+          <p style="margin-bottom: 10px; font-size: 94%">
+            This is a demo for Surface Normal ControlNet that using
+            <a href="https://huggingface.co/geyongtao/HumanWild" target="_blank"> HumanWild model</a> as backbone. 
+          </p>
+        ''')
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(sources=None, type="pil") # None for upload, ctrl+v and webcam
+                prompt = gr.Textbox(label="Prompt")
+                negative_prompt = gr.Textbox(label="Negative prompt", value="Logo,Watermark,Text,Ugly,Morbid,Extra fingers,Poorly drawn hands,Mutation,Blurry,Extra limbs,Gross proportions,Missing arms,Mutated hands,Long neck,Duplicate,Mutilated,Mutilated hands,Poorly drawn face,Deformed,Bad anatomy,Cloned face,Malformed limbs,Missing legs,Too many fingers")
+                num_steps = gr.Slider(label="Number of steps", minimum=25, maximum=100, value=50, step=1)
+                controlnet_conditioning_scale = gr.Slider(label="ControlNet conditioning scale", minimum=0.1, maximum=2.0, value=1.0, step=0.05)
+                seed = gr.Slider(label="Seed", minimum=0, maximum=2147483647, step=1, randomize=True,)
+                run_button = gr.Button(value="Run")
+                
+                
+            with gr.Column():
+                result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery", columns=[2], height='auto')
+        ips = [input_image, prompt, negative_prompt, num_steps, controlnet_conditioning_scale, seed]
+        
+        run_button.click(fn=process, inputs=ips, outputs=[result_gallery])
     
-    run_button.click(fn=process, inputs=ips, outputs=[result_gallery])
+    block.launch(debug = True)
 
-block.launch(debug = True)
+if __name__ == '__main__':
+    fire.Fire(run_demo)