Spaces:
Paused
Paused
File size: 2,956 Bytes
a0976c0 |
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 |
import gradio as gr
import os
from diffusers import DiffusionPipeline, FluxControlPipeline, FluxTransformer2DModel
import torch
from transformers import T5EncoderModel
from controlnet_aux import CannyDetector
from diffusers.utils import load_image
from huggingface_hub import login
hf_token = os.environ.get["HF_TOKEN"]
login(hf_token)
def load_pipeline(four_bit=False):
orig_pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
if four_bit:
print("Using four bit.")
transformer = FluxTransformer2DModel.from_pretrained(
"sayakpaul/FLUX.1-Canny-dev-nf4", subfolder="transformer", torch_dtype=torch.bfloat16
)
text_encoder_2 = T5EncoderModel.from_pretrained(
"sayakpaul/FLUX.1-Canny-dev-nf4", subfolder="text_encoder_2", torch_dtype=torch.bfloat16
)
pipeline = FluxControlPipeline.from_pipe(
orig_pipeline, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16
)
else:
transformer = FluxTransformer2DModel.from_pretrained(
"black-forest-labs/FLUX.1-Canny-dev",
subfolder="transformer",
revision="refs/pr/1",
torch_dtype=torch.bfloat16,
)
pipeline = FluxControlPipeline.from_pipe(orig_pipeline, transformer=transformer, torch_dtype=torch.bfloat16)
pipeline.enable_model_cpu_offload()
return pipeline
def get_canny(control_image):
processor = CannyDetector()
control_image = processor(
control_image, low_threshold=50, high_threshold=200, detect_resolution=1024, image_resolution=1024
)
return control_image
def main(ref_filepath, prompt, use_nf4):
pipe = load_pipeline(use_nf4)
control_image = load_image(ref_filepath)
control_image = get_canny(control_image)
image = pipe(
prompt=prompt,
control_image=control_image,
height=1024,
width=1024,
num_inference_steps=50,
guidance_scale=30.0,
max_sequence_length=512,
generator=torch.Generator("cpu").manual_seed(0),
).images[0]
filename = "output_"
filename += "_4bit" if four_bit else ""
image.save(f"{filename}.png")
return f"{filename}.png", control_image
with gr.Blocks() as demo:
with gr.Column():
gr.Markdown("# FLUX.1 Canny Dev")
with gr.Row():
with gr.Column():
image_input = gr.Image(label="Reference Image", type="filepath")
prompt = gr.Textbox(label="Prompt")
use_nf4 = gr.Checkbox(label="Use NF4 checkpoints", value=True)
submit_btn = gr.Button("Submit")
with gr.Column():
results= gr.Gallery(label="Results")
submit_btn.click(
fn = main,
inputs = [image_input, prompt, use_nf4],
outputs = [results]
)
demo.launch(show_api=False, show_error=True) |