ilanser commited on
Commit
7cdd9b7
·
1 Parent(s): 80cc72c

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +43 -0
app.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from PIL import Image
3
+ import base64
4
+ import io
5
+ import cv2
6
+ import numpy as np
7
+ import torch
8
+ from controlnet_aux import HEDdetector
9
+ from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
10
+
11
+ def predict(sketch, description):
12
+ # Convert sketch to PIL image
13
+ sketch_pil = Image.fromarray(sketch)
14
+
15
+ hed = HEDdetector.from_pretrained('lllyasviel/Annotators')
16
+
17
+ image = hed(sketch_pil, scribble=True)
18
+
19
+ model_id = "runwayml/stable-diffusion-v1-5"
20
+ controlnet_id = "lllyasviel/sd-controlnet-scribble"
21
+
22
+ # Load ControlNet model
23
+ controlnet = ControlNetModel.from_pretrained(controlnet_id, torch_dtype=torch.float16)
24
+
25
+ # Create pipeline with ControlNet model
26
+ pipe = StableDiffusionControlNetPipeline.from_pretrained(model_id, controlnet=controlnet, torch_dtype=torch.float16)
27
+
28
+ # Use improved scheduler
29
+ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
30
+
31
+ # Enable smart CPU offloading and memory efficient attention
32
+ # pipe.enable_model_cpu_offload()
33
+ pipe.enable_xformers_memory_efficient_attention()
34
+
35
+
36
+ result = pipe(description, image, num_inference_steps=20).images[0]
37
+
38
+ return result
39
+ # Define sketchpad with custom size and stroke width
40
+ sketchpad = gr.Sketchpad(shape=(1024, 1024), brush_radius=5)
41
+
42
+ iface = gr.Interface(fn=predict, inputs=[sketchpad, "text"], outputs="image", live=False)
43
+ iface.launch(share=True)