dy2000/optimized-diffusers-code
0
1import gradio as gr2from utils.pipeline_utils import determine_pipe_loading_memory3from utils.llm_utils import LLMCodeOptimizer4from prompts import system_prompt, generate_prompt5from utils.hardware_utils import categorize_ram, categorize_vram6 7LLM_CACHE = {}8 9 10def get_output_code(11 repo_id,12 gemini_model_to_use,13 disable_bf16,14 enbale_caching,15 enable_quantization,16 system_ram,17 gpu_vram,18 torch_compile_friendly,19 fp8_friendly,20 progress=gr.Progress(track_tqdm=True)21):22 loading_mem_out = determine_pipe_loading_memory(repo_id, None, disable_bf16)23 load_memory = loading_mem_out["total_loading_memory_gb"]24 ram_category = categorize_ram(system_ram)25 vram_category = categorize_vram(gpu_vram)26 27 print(f"RAM Category: {ram_category}")28 print(f"VRAM Category: {vram_category}")29 30 if gemini_model_to_use not in LLM_CACHE:31 print(f"Initializing new LLM instance for: {gemini_model_to_use}")32 # If not, create it and add it to the cache33 LLM_CACHE[gemini_model_to_use] = LLMCodeOptimizer(model_name=gemini_model_to_use, system_prompt=system_prompt)34 35 llm = LLM_CACHE[gemini_model_to_use]36 current_generate_prompt = generate_prompt.format(37 ckpt_id=repo_id,38 pipeline_loading_memory=load_memory,39 available_system_ram=system_ram,40 available_gpu_vram=gpu_vram,41 enable_caching=enable_caching,42 enable_quantization=enable_quantization,43 is_fp8_supported=fp8_friendly,44 enable_torch_compile=torch_compile_friendly,45 )46 generated_prompt = current_generate_prompt47 llm_output = llm(current_generate_prompt)48 return llm_output, generated_prompt49 50 51# --- Gradio UI Definition ---52with gr.Blocks() as demo:53 gr.Markdown(54 """55 # 🧨 Generate Diffusers Inference code snippet tailored to your machine56 Enter a Hugging Face Hub `repo_id` and your system specs to get started for inference.57 This tool uses [Gemini](https://ai.google.dev/gemini-api/docs/models) to generate the code based on your settings. This is based on58 [sayakpaul/auto-diffusers-docs](https://github.com/sayakpaul/auto-diffusers-docs/).59 """,60 elem_id="col-container"61 )62 63 with gr.Row():64 with gr.Column(scale=3):65 repo_id = gr.Textbox(66 label="Hugging Face Repo ID",67 placeholder="e.g., black-forest-labs/FLUX.1-dev",68 info="The model repository you want to analyze.",69 value="black-forest-labs/FLUX.1-dev",70 )71 gemini_model_to_use = gr.Dropdown(72 ["gemini-2.5-flash-lite", "gemini-2.5-flash", "gemini-2.5-pro"],73 value="gemini-2.5-flash-lite",74 label="Gemini Model",75 info="Select the model to generate the analysis.",76 )77 with gr.Row():78 system_ram = gr.Number(label="Free System RAM (GB)", value=20)79 gpu_vram = gr.Number(label="Free GPU VRAM (GB)", value=8)80 81 with gr.Row():82 disable_bf16 = gr.Checkbox(83 label="Disable BF16 (Use FP32)",84 value=False,85 info="Compute in 32-bit precision (caution ⚠️)",86 )87 enable_caching = gr.Checkbox(88 label="Enable lossy caching", value=False, info="Consider applying caching for speed"89 )90 enable_lossy = gr.Checkbox(91 label="Allow Lossy Quantization", value=False, info="Consider 8-bit/4-bit quantization"92 )93 torch_compile_friendly = gr.Checkbox(94 label="torch.compile() friendly", value=False, info="Model is compatible with torch.compile"95 )96 fp8_friendly = gr.Checkbox(97 label="fp8 friendly", value=False, info="Model and hardware support FP8 precision"98 )99 100 with gr.Column(scale=1):101 submit_btn = gr.Button("Get Code ☁", variant="primary", scale=1)102 103 # --- Start of New Code Block ---104 all_inputs = [105 repo_id,106 gemini_model_to_use,107 disable_bf16,108 enable_caching,109 enable_lossy,110 system_ram,111 gpu_vram,112 torch_compile_friendly,113 fp8_friendly,114 ]115 116 with gr.Accordion("Examples (Click to expand)", open=False):117 gr.Examples(118 examples=[119 [120 "stabilityai/stable-diffusion-xl-base-1.0",121 "gemini-2.5-pro",122 False,123 False,124 False,125 64,126 24,127 True,128 True,129 ],130 [131 "Wan-AI/Wan2.1-VACE-1.3B-diffusers",132 "gemini-2.5-flash",133 False,134 True,135 False,136 16,137 8,138 False,139 False,140 ],141 [142 "stabilityai/stable-diffusion-3-medium-diffusers",143 "gemini-2.5-pro",144 False,145 False,146 False,147 32,148 16,149 True,150 False,151 ],152 ],153 inputs=all_inputs,154 label="Examples (Click to try)",155 )156 # --- End of New Code Block ---157 158 with gr.Accordion("💡 Tips", open=False):159 gr.Markdown(160 """161 - Try changing to the model from Flash to Pro if the results are bad.162 - Please provide the VRAM and RAM details accurately as the suggestions depend on them.163 - As a rule of thumb, GPUs from RTX 4090 and later, are generally good for using `torch.compile()`.164 - When lossy quantization isn't preferred try enabling caching. Caching can still be lossy, though.165 - To leverage FP8, the GPU needs to have a compute capability of at least 8.9.166 - Check out the following docs for optimization in Diffusers:167 * [Memory](https://huggingface.co/docs/diffusers/main/en/optimization/memory)168 * [Caching](https://huggingface.co/docs/diffusers/main/en/optimization/cache)169 * [Inference acceleration](https://huggingface.co/docs/diffusers/main/en/optimization/fp16)170 * [PyTorch blog](https://pytorch.org/blog/presenting-flux-fast-making-flux-go-brrr-on-h100s/)171 """172 )173 174 with gr.Accordion("Generated LLM Prompt (for debugging)", open=False):175 prompt_output = gr.Textbox(label="Prompt", show_copy_button=True, lines=10, interactive=False)176 177 gr.Markdown("---")178 179 with gr.Accordion("Generated Code 💻", open=True):180 code_output = gr.Code(interactive=True, language="python")181 182 gr.Markdown(183 """184 ---185 > ⛔️ **Disclaimer:** Large Language Models (LLMs) can make mistakes. The information provided186 > is an estimate and should be verified. Always test the model on your target hardware to confirm187 > actual memory requirements.188 """189 )190 191 # --- Event Handling ---192 submit_btn.click(fn=get_output_code, inputs=all_inputs, outputs=[code_output, prompt_output])193 194 195if __name__ == "__main__":196 demo.launch()197