| import spaces
|
| import gradio as gr
|
| import torch
|
| from huggingface_hub import hf_hub_download
|
| from diffusers import FluxPipeline, FluxTransformer2DModel, GGUFQuantizationConfig, BitsAndBytesConfig
|
| import os
|
|
|
|
|
|
|
|
|
|
|
| HF_TOKEN = os.getenv("HF_TOKEN", "")
|
| device = "cuda" if torch.cuda.is_available() else "cpu"
|
| flux_repo = "multimodalart/FLUX.1-dev2pro-full"
|
| ckpt_path = "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf"
|
| transformer_gguf = FluxTransformer2DModel.from_single_file(ckpt_path, subfolder="transformer", quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
|
| torch_dtype=torch.bfloat16, config=flux_repo, token=HF_TOKEN)
|
| transformer = FluxTransformer2DModel.from_pretrained(flux_repo, subfolder="transformer", torch_dtype=torch.bfloat16, token=HF_TOKEN)
|
| nf4_quantization_config = BitsAndBytesConfig(load_in_4bit=True)
|
| transformer_nf4 = FluxTransformer2DModel.from_pretrained(flux_repo, subfolder="transformer", quantization_config=nf4_quantization_config,
|
| torch_dtype=torch.bfloat16, token=HF_TOKEN)
|
| pipe = FluxPipeline.from_pretrained(flux_repo, transformer=transformer, torch_dtype=torch.bfloat16, token=HF_TOKEN)
|
| hyper_sd_lora = hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors")
|
|
|
| @spaces.GPU(duration=70)
|
| def infer(prompt: str, mode: str, is_lora: bool, progress=gr.Progress(track_tqdm=True)):
|
| global pipe
|
| try:
|
| pipe.unload_lora_weights()
|
| if mode == "Default": pipe.transformer = transformer
|
| elif mode == "GGUF": pipe.transformer = transformer_gguf
|
| elif mode == "NF4": pipe.transformer = transformer_nf4
|
| if is_lora:
|
| pipe.load_lora_weights(hyper_sd_lora, adapter_name="hyper-sd")
|
| pipe.set_adapters(["hyper-sd"], adapter_weights=[0.125])
|
| steps = 8
|
| else: steps = 28
|
| pipe.to(device)
|
| image = pipe(prompt, generator=torch.manual_seed(0), num_inference_steps=steps).images[0]
|
| pipe.to("cpu")
|
| return image
|
| except Exception as e:
|
| raise gr.Error(e)
|
|
|
| with gr.Blocks() as demo:
|
| with gr.Row():
|
| with gr.Column():
|
| prompt = gr.Textbox(label="Prompt", value="A cat holding a sign that says hello world", lines=1)
|
| mode = gr.Radio(label="Mode", choices=["Default", "GGUF", "NF4"], value="Default")
|
| is_lora = gr.Checkbox(label="Enable LoRA", value=True)
|
| gen_btn = gr.Button("Generate Image")
|
| with gr.Column():
|
| result = gr.Image(label="Result Image")
|
|
|
| gen_btn.click(infer, [prompt, mode, is_lora], [result])
|
|
|
| demo.launch()
|
|
|
