|
|
import gradio as gr |
|
|
from diffusers import StableDiffusionPipeline, StableDiffusionDepth2ImgPipeline |
|
|
from PIL import Image |
|
|
|
|
|
|
|
|
sd_pipeline = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1") |
|
|
sd_pipeline.to("cuda") |
|
|
depth_pipeline = StableDiffusionDepth2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-depth2img") |
|
|
depth_pipeline.to("cuda") |
|
|
|
|
|
def generate_images(prompt): |
|
|
base_image = sd_pipeline(prompt).images[0] |
|
|
depth_image = depth_pipeline(prompt=prompt, image=base_image).images[0] |
|
|
return base_image, depth_image |
|
|
|
|
|
iface = gr.Interface( |
|
|
fn=generate_images, |
|
|
inputs=gr.Textbox(label="Prompt"), |
|
|
outputs=[gr.Image(label="Bas-relief"), gr.Image(label="Depth Map")], |
|
|
title="Text to Bas-Relief and Depth Map", |
|
|
) |
|
|
|
|
|
iface.launch() |
|
|
|
