Spaces:
Runtime error
Runtime error
File size: 2,875 Bytes
dd7bda3 8ae18c8 dd7bda3 8ae18c8 dd7bda3 8ae18c8 dd7bda3 8ae18c8 306edbc dd7bda3 1efb400 dd7bda3 1efb400 dd7bda3 8ae18c8 dd7bda3 8ae18c8 |
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 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import gradio as gr
import torch
import spaces
from diffusers import FluxPipeline, DiffusionPipeline
from torchao.quantization import autoquant
# # # normal FluxPipeline
# pipeline_normal = FluxPipeline.from_pretrained(
# "sayakpaul/FLUX.1-merged",
# torch_dtype=torch.bfloat16
# ).to("cuda")
# pipeline_normal.transformer.to(memory_format=torch.channels_last)
# pipeline_normal.transformer = torch.compile(pipeline_normal.transformer, mode="max-autotune", fullgraph=True)
pipeline_normal = DiffusionPipeline.from_pretrained("sayakpaul/FLUX.1-merged")
pipeline_normal.load_lora_weights("DarkMoonDragon/TurboRender-flux-dev")
# # optimized FluxPipeline
# pipeline_optimized = FluxPipeline.from_pretrained(
# "camenduru/FLUX.1-dev-diffusers",
# torch_dtype=torch.bfloat16
# ).to("cuda")
# pipeline_optimized.transformer.to(memory_format=torch.channels_last)
# pipeline_optimized.transformer = torch.compile(
# pipeline_optimized.transformer,
# mode="max-autotune",
# fullgraph=True
# )
# # wrap the autoquant call in a try-except block to handle unsupported layers
# for name, layer in pipeline_optimized.transformer.named_children():
# try:
# # apply autoquant to each layer
# pipeline_optimized.transformer._modules[name] = autoquant(layer, error_on_unseen=False)
# print(f"Successfully quantized {name}")
# except AttributeError as e:
# print(f"Skipping layer {name} due to error: {e}")
# except Exception as e:
# print(f"Unexpected error while quantizing {name}: {e}")
# pipeline_optimized.transformer = autoquant(
# pipeline_optimized.transformer,
# error_on_unseen=False
# )
pipeline_optimized = pipeline_normal
@spaces.GPU(duration=120)
def generate_images(prompt, guidance_scale, num_inference_steps):
# # generate image with normal pipeline
# image_normal = pipeline_normal(
# prompt=prompt,
# guidance_scale=guidance_scale,
# num_inference_steps=int(num_inference_steps)
# ).images[0]
# generate image with optimized pipeline
image_optimized = pipeline_optimized(
prompt=prompt,
guidance_scale=guidance_scale,
num_inference_steps=int(num_inference_steps)
).images[0]
return image_optimized
# set up Gradio interface
demo = gr.Interface(
fn=generate_images,
inputs=[
gr.Textbox(lines=2, placeholder="Enter your prompt here...", label="Prompt"),
gr.Slider(1.0, 10.0, step=0.5, value=3.5, label="Guidance Scale"),
gr.Slider(10, 100, step=1, value=50, label="Number of Inference Steps")
],
outputs=[
gr.Image(type="pil", label="Optimized FluxPipeline")
],
title="FluxPipeline Comparison",
description="Compare images generated by the normal FluxPipeline and the optimized one using torchao and torch.compile()."
)
demo.launch()
|