File size: 1,944 Bytes
0fd6850 |
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 |
import os
import torch
from PIL import Image
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
import gradio as gr
# Disable oneDNN custom operations
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
# Clear PyTorch cache
torch.cuda.empty_cache()
# Check if CUDA is available
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cuda":
print("CUDA is available. Device count:", torch.cuda.device_count())
print("Current device:", torch.cuda.current_device())
print("Device name:", torch.cuda.get_device_name(torch.cuda.current_device()))
else:
print("CUDA is not available. Using CPU.")
# Load ControlNet model with OpenPose pre-trained weights from Hugging Face
controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_openpose", torch_dtype=torch.float16)
# Load the Stable Diffusion model
pipe = StableDiffusionControlNetPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
).to(device)
# Function for inference
def generate_image(prompt, target_image, pose_image):
try:
# Resize images
target_image = target_image.resize((512, 512))
pose_image = pose_image.resize((512, 512))
# Generate image with ControlNet
output = pipe(prompt=prompt, image=target_image, control_image=pose_image, num_inference_steps=50)
# Return the result
return output["sample"][0]
except Exception as e:
print(f"Error during image generation: {e}")
return None
# Setup Gradio Interface
interface = gr.Interface(
fn=generate_image,
inputs=[
gr.Textbox(label="Prompt"),
gr.Image(label="Target Image", type="pil"),
gr.Image(label="Pose Image (Reference)", type="pil")
],
outputs=gr.Image(label="Generated Image")
)
# Launch the interface
interface.launch() |