Spaces:
Runtime error
Runtime error
File size: 2,945 Bytes
89b3db2 754b60e 1368e65 754b60e e19c312 754b60e e19c312 754b60e e19c312 754b60e 1368e65 213883e 1368e65 92e2f62 14d5805 92e2f62 14d5805 92e2f62 14d5805 9c0158e d1be92b 213883e 1368e65 e19c312 1368e65 ba6d92c 14d5805 1368e65 ba6d92c 1368e65 e19c312 |
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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
import torch
from diffusers.utils import load_image
from diffusers import FluxControlNetPipeline, FluxControlNetModel, FluxMultiControlNetModel
import gradio as gr
import spaces
# Ensure that you're using the appropriate data type for your GPU
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
base_model = 'black-forest-labs/FLUX.1-dev'
controlnet_model_union = 'InstantX/FLUX.1-dev-Controlnet-Union'
controlnet_union = FluxControlNetModel.from_pretrained(controlnet_model_union, torch_dtype=torch_dtype)
controlnet = FluxMultiControlNetModel([controlnet_union])
pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch_dtype)
# If you encounter issues with CUDA, you can run this on the CPU for debugging
pipe.to("cuda" if torch.cuda.is_available() else "cpu")
control_modes = [
"canny",
"tile",
"depth",
"blur",
"pose",
"gray",
"lq"
]
def adjust_dimensions(width, height):
adjusted_width = width - (width % 8)
adjusted_height = height - (height % 8)
return adjusted_width, adjusted_height
@spaces.GPU
def generate_image(prompt, control_image_depth, control_mode_depth_index, use_depth, control_image_canny, control_mode_canny_index):
control_images = []
control_modes = []
conditioning_scales = []
if use_depth:
control_images.append(control_image_depth)
control_modes.append(control_mode_depth_index)
conditioning_scales.append(0.2)
control_images.append(control_image_canny)
control_modes.append(control_mode_canny_index)
conditioning_scales.append(0.4)
width, height = control_image_canny.shape[:2]
adjusted_width, adjusted_height = adjust_dimensions(width, height)
try:
image = pipe(
prompt,
control_image=control_images,
control_mode=control_modes,
width=adjusted_width,
height=adjusted_height,
controlnet_conditioning_scale=conditioning_scales,
num_inference_steps=24,
guidance_scale=3.5,
generator=torch.manual_seed(42),
).images[0]
except RuntimeError as e:
torch.cuda.empty_cache()
raise e
return image
iface = gr.Interface(
fn=generate_image,
inputs=[
gr.Text(label="Prompt"),
gr.Image(label="Control Image (Depth)"),
gr.Dropdown(choices=control_modes, value=control_modes.index("depth"), label="Control Mode (Depth)"),
gr.Checkbox(label="Use Depth Control Image", value=True),
gr.Image(label="Control Image (Canny)"),
gr.Dropdown(choices=control_modes, value=control_modes.index("canny"), label="Control Mode (Canny)")
],
outputs=gr.Image(label="Generated Image"),
title="FluxControlNet Image Generation",
description="Generate an image using FluxControlNet with depth and canny control images.",
)
iface.launch(share=True)
|