File size: 5,961 Bytes
2414743 |
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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 |
import random
import torch
from PIL import Image
import gradio as gr
from diffusers import DiffusionPipeline
# Configure deterministic behavior for reproducibility
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.backends.cuda.matmul.allow_tf32 = True
MAX_SEED = 2**32 - 1
class ModelManager:
"""
Handles model initialization, LoRA weight loading, and image generation.
"""
def __init__(self, base_model: str, lora_repo: str, trigger_word: str = ""):
self.trigger_word = trigger_word
self.pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
self.pipe.load_lora_weights(lora_repo)
self.pipe.to("cuda")
def generate_image(self, prompt: str, cfg_scale: float, steps: int, seed: int,
width: int, height: int, lora_scale: float, progress_callback) -> Image.Image:
"""
Generates an image based on the given prompt and parameters using a callback for progress updates.
"""
# Establish reproducible generator
generator = torch.Generator(device="cuda").manual_seed(seed)
full_prompt = f"{prompt} {self.trigger_word}"
def callback_fn(step: int, timestep: int, latents):
percentage = int((step / steps) * 100)
message = f"Processing step {step} of {steps}..."
progress_callback(percentage, message)
# Generate image with integrated progress reporting
image = self.pipe(
prompt=full_prompt,
num_inference_steps=steps,
guidance_scale=cfg_scale,
width=width,
height=height,
generator=generator,
joint_attention_kwargs={"scale": lora_scale},
callback=callback_fn,
callback_steps=1,
).images[0]
return image
# Initialize the model manager with specified models and LoRA weights
model_manager = ModelManager(
base_model="black-forest-labs/FLUX.1-dev",
lora_repo="XLabs-AI/flux-RealismLora",
trigger_word=""
)
def run_generation(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
"""
Gradio interface callback to manage seed randomization, progress updates,
and image generation using the ModelManager.
"""
if randomize_seed:
seed = random.randint(0, MAX_SEED)
# Start the progress
progress(0, "Starting image generation...")
# Generate the image using the model manager with progress callback integration
image = model_manager.generate_image(
prompt, cfg_scale, steps, seed, width, height, lora_scale, progress
)
# Mark completion
progress(100, "Completed!")
return image, seed
# Example parameters and image path for initializing the interface with defaults
example_image_path = "example0.webp"
example_prompt = (
"A Jelita Sukawati speaker is captured mid-speech. She has long, dark brown hair that cascades over her shoulders, "
"framing her radiant, smiling face. Her Latina features are highlighted by warm, sun-kissed skin and bright, "
"expressive eyes. She gestures with her left hand, displaying a delicate ring on her pinky finger, as she speaks passionately. "
"The woman is wearing a colorful, patterned dress with a green lanyard featuring multiple badges and logos hanging around her neck. "
"The lanyard prominently displays the 'CagliostroLab' text. Behind her, there is a blurred background with a white banner "
"containing logos and text, indicating a professional or conference setting. The overall scene captures the energy and vibrancy "
"of her presentation."
)
example_cfg_scale = 3.2
example_steps = 32
example_width = 1152
example_height = 896
example_seed = 3981632454
example_lora_scale = 0.85
def load_example():
# Load example image for initial display
example_image = Image.open(example_image_path)
return (
example_prompt,
example_cfg_scale,
example_steps,
True,
example_seed,
example_width,
example_height,
example_lora_scale,
example_image
)
with gr.Blocks() as app:
gr.Markdown("# Flux RealismLora Image Generator")
with gr.Row():
with gr.Column(scale=3):
prompt = gr.TextArea(label="Prompt", placeholder="Type a prompt", lines=5)
generate_button = gr.Button("Generate")
cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=example_cfg_scale)
steps = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=example_steps)
width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=example_width)
height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=example_height)
randomize_seed = gr.Checkbox(True, label="Randomize seed")
seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=example_seed)
lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=1, step=0.01, value=example_lora_scale)
with gr.Column(scale=1):
result = gr.Image(label="Generated Image")
gr.Markdown(
"Generate images using RealismLora and a text prompt.\n"
"[[non-commercial license, Flux.1 Dev](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)]"
)
# Load example data on launch
app.load(
load_example,
inputs=[],
outputs=[prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, result]
)
# Set up button interaction
generate_button.click(
run_generation,
inputs=[prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale],
outputs=[result, seed]
)
app.queue()
app.launch()
|