|
|
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()
|
|
|
|
