Spaces:
Runtime error
Runtime error
| 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 | |
| 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) | |