import spaces import argparse import os import time from os import path from safetensors.torch import load_file from huggingface_hub import hf_hub_download cache_path = path.join(path.dirname(path.abspath(__file__)), "models") os.environ["TRANSFORMERS_CACHE"] = cache_path os.environ["HF_HUB_CACHE"] = cache_path os.environ["HF_HOME"] = cache_path import gradio as gr import torch from diffusers import FluxPipeline torch.backends.cuda.matmul.allow_tf32 = True loaded_acc = None class timer: def __init__(self, method_name="timed process"): self.method = method_name def __enter__(self): self.start = time.time() print(f"{self.method} starts") def __exit__(self, exc_type, exc_val, exc_tb): end = time.time() print(f"{self.method} took {str(round(end - self.start, 2))}s") if not path.exists(cache_path): os.makedirs(cache_path, exist_ok=True) pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16) pipe.load_lora_weights(hf_hub_download("RED-AIGC/TDD", "FLUX.1-dev_tdd_lora_weights.safetensors"),adapter_name="TDD") pipe.load_lora_weights(hf_hub_download("RED-AIGC/TDD", "FLUX.1-dev_tdd_adv_lora_weights.safetensors"),adapter_name="TDD_adv") # pipe.fuse_lora(lora_scale=0.125) pipe.to("cuda") css = """ h1 { text-align: center; display:block; } .gradio-container { max-width: 70.5rem !important; } """ @spaces.GPU def process_image(prompt,acc,height, width, steps, scales, seed): global pipe global loaded_acc if loaded_acc != acc: #pipe.load_lora_weights(ACC_lora[acc], adapter_name=acc) pipe.set_adapters([acc], adapter_weights=[0.125]) print(pipe.get_active_adapters()) loaded_acc = acc with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"): return pipe( prompt=[prompt], generator=torch.Generator().manual_seed(int(seed)), num_inference_steps=int(steps), guidance_scale=float(scales), height=int(height), width=int(width), max_sequence_length=256 ).images[0] with gr.Blocks(css=css) as demo: gr.Markdown( """ # FLUX.1-dev(beta) distilled by ✨Target-Driven Distillation✨ Compared to Hyper-FLUX, the beta version of TDD has its parameters reduced by half(600MB), resulting in more realistic details. Due to limitations in machine resources, there are still many imperfections in the beta version. Besides, TDD is also available for distilling video generation models. This space presents TDD-distilled [FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev). [**Project Page**](https://redaigc.github.io/TDD/) **|** [**Paper**](https://arxiv.org/abs/2409.01347) **|** [**Code**](https://github.com/RedAIGC/Target-Driven-Distillation) **|** [**Model**](https://huggingface.co/RED-AIGC/TDD) **|** [🤗 **TDD-SDXL Demo**](https://huggingface.co/spaces/RED-AIGC/TDD) **|** [🤗 **TDD-SVD Demo**](https://huggingface.co/spaces/RED-AIGC/SVD-TDD) The codes of this space are built on [Hyper-FLUX](https://huggingface.co/spaces/ByteDance/Hyper-FLUX-8Steps-LoRA) and we acknowledge their contribution. """ ) with gr.Row(): with gr.Column(scale=3): with gr.Group(): prompt = gr.Textbox( label="Prompt", # value="portrait photo of a girl, photograph, highly detailed face, depth of field, moody light, golden hour, style by Dan Winters, Russell James, Steve McCurry, centered, extremely detailed, Nikon D850, award winning photography", lines=3 ) with gr.Accordion("Advanced Settings", open=False): with gr.Group(): with gr.Row(): height = gr.Slider(label="Height", minimum=256, maximum=1152, step=64, value=1024) width = gr.Slider(label="Width", minimum=256, maximum=1152, step=64, value=1024) with gr.Row(): steps = gr.Slider(label="Inference Steps", minimum=4, maximum=10, step=1, value=8) scales = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=5.0, step=0.1, value=2.0) with gr.Row(): seed = gr.Number(label="Seed", value=-1, precision=0) with gr.Row(): acc = gr.Dropdown( label="Accelerate Lora", choices=["TDD", "TDD_adv"], value="TDD_adv", ) generate_btn = gr.Button("Generate Image", variant="primary", scale=1) with gr.Column(scale=4): output = gr.Image(label="Your Generated Image") person1="A young woman with strikingly symmetrical features, smooth porcelain skin, and large hazel eyes framed by thick, wavy auburn hair, sitting beside a window where soft sunlight highlights her high cheekbones and glossy lips." person2="A radiant woman with soft, youthful skin, her lips lightly tinted in a natural pink, and her large, round eyes framed by thick lashes, posing gracefully in a traditional kimono among red maple leaves." person3="A modern woman with striking blue eyes, a perfectly symmetrical face, and soft, glossy lips, sitting in a sleek interior with soft, natural light highlighting her flawless skin." dog="Portrait photo of a Shiba Inu, photograph, highly detailed fur, warm and cheerful light, soft pastel tones, vibrant and sunny atmosphere, style by Tim Flach, bright daylight, natural setting, centered, extremely detailed, Nikon D850, award-winning photography" fox="A majestic Arctic fox standing gracefully on a snowy tundra, its thick white fur blending seamlessly with the icy surroundings. The texture of the fox's fur is highly detailed, showcasing its fluffy and warm coat. " scenery1="A quiet autumn street lined with maple trees, their vibrant red and orange leaves falling gently to the ground. The pavement is partially covered with scattered leaves, while a soft breeze stirs the branches. Sunlight filters through the canopy, casting warm, dappled light on the sidewalk. A lone cyclist rides down the street, and a few pedestrians stroll along, enjoying the crisp autumn air. The scene captures the serene beauty and subtle movement of an autumn day." scenery2="A high-detail close-up of the Forbidden City's architectural details, with snowflakes gracefully falling. Shot on a Nikon D850, the 45.7-megapixel sensor captures the scene's bright warmth. A fast shutter speed at a low ISO freezes the snowflakes, ensuring a crisp, vivid image. Ideal for a winter magazine feature." gr.Examples( examples=[ [person1, "TDD_adv", 1024, 1024, 8, 2.1, 5685], [person2, "TDD_adv", 1024, 1024, 8, 2, 8888], [person3,"TDD_adv", 1024, 1024, 6, 1.8, 3420], [dog, "TDD", 1024, 1024, 6, 2, 29], [fox,"TDD_adv", 1024, 1024, 8, 2.1, 4678], [scenery1, "TDD_adv", 1024, 1024, 8, 2.5, 9669], [scenery2, "TDD_adv", 1024, 1024, 6, 1.8, 3420], ], # inputs=[prompt, negative_prompt, ckpt, acc, steps, guidance_scale, eta, seed], inputs=[prompt,acc, height, width, steps, scales, seed], outputs=output, fn=process_image, cache_examples="lazy", ) gr.Markdown( """

How to Use

  1. Enter a detailed description of the image you want to create.
  2. Adjust advanced settings if desired (tap to expand).
  3. Tap "Generate Image" and wait for your creation!

Tip: Be specific in your description for best results!

""" ) generate_btn.click( process_image, inputs=[prompt, acc,height, width, steps, scales, seed], outputs=output ) if __name__ == "__main__": demo.launch()