File size: 1,928 Bytes
86a2bd4
bb21a7c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e5d785e
 
 
bb21a7c
 
 
 
 
 
 
 
 
 
 
 
 
e5d785e
 
bb21a7c
a02cd43
 
1b8870f
 
 
bb21a7c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
from controlnet_aux import OpenposeDetector
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers import UniPCMultistepScheduler
import gradio as gr
import torch

# Constants
low_threshold = 100
high_threshold = 200

# Models
pose_model = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
controlnet = ControlNetModel.from_pretrained(
    "lllyasviel/sd-controlnet-openpose", torch_dtype=torch.float16
)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "andite/anything-v4.0", controlnet=controlnet, torch_dtype=torch.float16
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)

# This command loads the individual model components on GPU on-demand. So, we don't
# need to explicitly call pipe.to("cuda").
pipe.enable_model_cpu_offload()

# xformers
pipe.enable_xformers_memory_efficient_attention()

# Generator seed,
generator = torch.manual_seed(0)

def get_pose(image):
    return pose_model(image) 


def generate_images(image, prompt):
    pose = get_pose(image)
    output = pipe(
        prompt,
        pose,
        generator=generator,
        num_images_per_prompt=3,
        num_inference_steps=20,
    )
    all_outputs = []
    all_outputs.append(pose)
    for image in output.images:
        all_outputs.append(image)
    return all_outputs


gr.Interface(
    generate_images,
    inputs=[
        gr.Image(type="pil"),
        gr.Textbox(
            label="Enter your prompt",
            max_lines=1,
            placeholder="best quality, extremely detailed",
        ),
    ],
    outputs=gr.Gallery().style(grid=[2], height="auto"),
    title="Generate controlled outputs with ControlNet and Stable Diffusion. ",
    description="This Space uses pose estimated lines as the additional conditioning.",
    examples=[["yoga1.jpeg", "best quality, extremely detailed"]],
    allow_flagging=False,
).launch(enable_queue=True)