Spaces:
Runtime error
Runtime error
File size: 2,935 Bytes
9feb242 |
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 |
import spaces
import gradio as gr
import torch
from huggingface_hub import hf_hub_download
from diffusers import FluxPipeline, FluxTransformer2DModel, GGUFQuantizationConfig, BitsAndBytesConfig
import os
import subprocess
#subprocess.run("pip list", shell=True)
#subprocess.run("diffusers-cli env", shell=True)
#from optimum.quanto import freeze, qfloat8, quantize
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()
|