DjStompzone/ControlNet-Flux-1-LineArt

1import gradio as gr2import torch3from diffusers.utils import load_image4from diffusers.pipelines.flux.pipeline_flux_controlnet import FluxControlNetPipeline5from diffusers.models.controlnet_flux import FluxControlNetModel6import random7import numpy as np8 9import os10from huggingface_hub import login11 12login(os.getenv("hfapikey"))13 14# Initialize models15base_model = 'black-forest-labs/FLUX.1-dev'16controlnet_model = 'promeai/FLUX.1-controlnet-lineart-promeai'17device = "cuda" if torch.cuda.is_available() else "cpu"18torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float3219 20controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch_dtype)21pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch_dtype)22pipe = pipe.to(device)23 24MAX_SEED = np.iinfo(np.int32).max25 26def infer(27    prompt,28    control_image_path,29    controlnet_conditioning_scale,30    guidance_scale,31    num_inference_steps,32    seed,33    randomize_seed,34):35    if randomize_seed:36        seed = random.randint(0, MAX_SEED)37 38    generator = torch.manual_seed(seed)39    control_image = load_image(control_image_path) if control_image_path else None40 41    # Generate image42    result = pipe(43        prompt=prompt,44        control_image=control_image,45        controlnet_conditioning_scale=controlnet_conditioning_scale,46        num_inference_steps=num_inference_steps,47        guidance_scale=guidance_scale,48        generator=generator,49    ).images[0]50 51    return result, seed52 53css = """54#col-container {55    margin: 0 auto;56    max-width: 640px;57}58"""59 60with gr.Blocks(css=css) as demo:61    with gr.Column(elem_id="col-container"):62        gr.Markdown("## Zero-shot Partial Style Transfer for Line Art Images, Powered by FLUX.1")63 64        with gr.Row():65            prompt = gr.Textbox(66                label="Prompt",67                placeholder="Enter your prompt",68                max_lines=1,69            )70            run_button = gr.Button("Generate", variant="primary")71 72        result = gr.Image(label="Result", show_label=False)73 74        with gr.Accordion("Advanced Settings", open=False):75            control_image = gr.Image(76                source="upload",77                type="filepath",78                label="Control Image (Line Art)"79            )80            controlnet_conditioning_scale = gr.Slider(81                label="ControlNet Conditioning Scale",82                minimum=0.0,83                maximum=1.0,84                value=0.6,85                step=0.186            )87            guidance_scale = gr.Slider(88                label="Guidance Scale",89                minimum=1.0,90                maximum=10.0,91                value=3.5,92                step=0.193            )94            num_inference_steps = gr.Slider(95                label="Number of Inference Steps",96                minimum=1,97                maximum=100,98                value=28,99                step=1100            )101            seed = gr.Slider(102                label="Seed",103                minimum=0,104                maximum=MAX_SEED,105                step=1,106                value=0107            )108            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)109 110        gr.Examples(111            examples=[112                "Anime girl with fennec ears holding a cake",113                "Victorian style mansion interior with candlelight"114            ],115            inputs=[prompt]116        )117 118    run_button.click(119        infer,120        inputs=[121            prompt,122            control_image,123            controlnet_conditioning_scale,124            guidance_scale,125            num_inference_steps,126            seed,127            randomize_seed128        ],129        outputs=[result, seed]130    )131 132if __name__ == "__main__":133    demo.launch()