Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,752 Bytes
eb1e114 366a699 eb1e114 366a699 eb1e114 2d87b1e eb1e114 |
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 |
import torch
import spaces
import os
from diffusers.utils import load_image
from diffusers import FluxControlNetModel, FluxControlNetPipeline, AutoencoderKL
import gradio as gr
huggingface_token = os.getenv("HUGGINFACE_TOKEN")
good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16, use_safetensors=True, token=huggingface_token).to("cuda")
# Load pipeline
controlnet = FluxControlNetModel.from_pretrained(
"jasperai/Flux.1-dev-Controlnet-Upscaler",
torch_dtype=torch.bfloat16
)
pipe = FluxControlNetPipeline.from_pretrained(
"LPX55/FLUX.1-merged_uncensored",
controlnet=controlnet,
torch_dtype=torch.bfloat16,
device_map="auto",
vae=good_vae,
use_safetensors=True,
token=huggingface_token
)
# Add to your pipeline initialization:
pipe.enable_xformers_memory_efficient_attention()
# pipe.enable_vae_slicing() # Batch processing of VAE
pipe.enable_model_cpu_offload() # Use with accelerate
# Convert all models to memory-efficient format
pipe.to(memory_format=torch.channels_last)
pipe.to("cuda")
@spaces.GPU
def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale):
# Load control image
control_image = control_image.resize((int(w * scale), int(h * scale)), PIL.Image.BICUBIC)
# control_image = load_image(control_image)
w, h = control_image.size
# Upscale x1
control_image = control_image.resize((int(w * scale), int(h * scale)))
print("Size to: " + str(control_image.size[0]) + ", " + str(control_image.size[1]))
image = pipe(
prompt=prompt,
control_image=control_image,
controlnet_conditioning_scale=controlnet_conditioning_scale,
num_inference_steps=steps,
guidance_scale=guidance_scale,
height=control_image.size[1],
width=control_image.size[0],
torch_dtype=torch.bfloat16,
device_map="auto"
).images[0]
torch.cuda.empty_cache()
return image
# Create Gradio interface
iface = gr.Interface(
fn=generate_image,
inputs=[
gr.Textbox(lines=2, placeholder="Enter your prompt here..."),
gr.Slider(1, 3, value=1, label="Scale"),
gr.Slider(6, 30, value=8, label="Steps"),
gr.Image(type="pil", label="Control Image"),
gr.Slider(0, 1, value=0.6, label="ControlNet Scale"),
gr.Slider(1, 20, value=3.5, label="Guidance Scale"),
],
outputs=[
gr.Image(type="pil", label="Generated Image", format="png"),
],
title="FLUX ControlNet Image Generation",
description="Generate images using the FluxControlNetPipeline. Upload a control image and enter a prompt to create an image.",
)
# Launch the app
iface.launch() |