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.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=[1.0]) 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", ) generate_btn = gr.Button("Generate Image", variant="primary", scale=1) with gr.Column(scale=4): output = gr.Image(label="Your Generated Image") person="Portrait of a girl, high-end photography, prominent facial details, significant depth of field, soft lighting, late afternoon sun, in the style of Richard Avedon, Henri Cartier-Bresson, and Yousuf Karsh. Golden ratio composition, exquisite detail, Hasselblad H6D-100c, featured on the cover of photography magazines." person2="A high-end, close-up portrait of a young woman and a man outdoors in soft, golden late afternoon sunlight. The woman is wearing an orange jacket, and the man is wearing a black jacket. The focus is on their faces, featuring ultra-high-resolution details such as skin texture, natural facial expressions, and realistic lighting. The image has a shallow depth of field to emphasize the sharpness of their facial features while softly blurring the background. Captured in the style of professional magazine photography with golden ratio composition, exquisite clarity, and a modern touch. The atmosphere is warm, serene, and natural." person1="Snowy Lake Portrait, fine-art photography, capturing a woman's serene expression with intricate facial features, against a backdrop of a frosted lake at twilight. The shallow depth of field isolates the subject from the softly lit, ambient snowscape. Inspired by the styles of Annie Leibovitz, Joe McNally, and David LaChapelle, this centered composition is extremely detailed, shot on a Nikon D850, and worthy of a photography magazine cover." 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" bird="Seize the moment with a high-detail shot of a crane's majestic flight using a Nikon D850. The 45.7-megapixel sensor ensures sharp feather detail. Paired with a Nikkor 600mm f/4E FL ED VR lens, the setup captures the bird with a fast shutter speed against a serene wetland backdrop. Low ISO maintains image clarity, and a shallow depth of field isolates the crane, making it a standout in wildlife photography." scenery1="Capture the golden hour's warm embrace on a sandy beach with a Nikon D850, highlighting the vibrant blues and the soft, golden sunlight. This award-winning photography captures the essence of a summer sunset by the coast." 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=[ [person, "TDD", 1024, 1024, 8, 2, 8888], [person1, "TDD", 1024, 1024, 6, 2, 2929], [person2, "TDD", 1024, 1024, 8, 2, 8888], [dog, "TDD", 1024, 1024, 6, 2, 29], [bird,"TDD", 1024, 1024, 8, 2.2, 7777], [scenery2, "TDD", 1024, 1024, 6, 1.8, 3420], [scenery1, "TDD", 1024, 1024, 6, 2.2, 3777], ], # 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()