Spaces:
Running
on
Zero
Running
on
Zero
File size: 8,737 Bytes
7f4c99b aa0cb15 7f4c99b fc39399 7f4c99b 3a26d89 75ba08b 7f4c99b 75ba08b 7f4c99b 75ba08b 7f4c99b 2b7d18b 7f4c99b |
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 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 |
import os
import subprocess
import sys
import io
import gradio as gr
import numpy as np
import random
import spaces
import torch
from diffusers import Flux2Pipeline, Flux2Transformer2DModel
from diffusers import BitsAndBytesConfig as DiffBitsAndBytesConfig
from optimization import optimize_pipeline_
import requests
from PIL import Image
import json
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024
def remote_text_encoder(prompts):
response = requests.post(
"https://remote-text-encoder-flux-2.huggingface.co/predict",
json={"prompt": prompts},
headers={
"Authorization": f"Bearer {os.environ['HF_TOKEN']}",
"Content-Type": "application/json"
}
)
assert response.status_code == 200, f"{response.status_code=}"
prompt_embeds = torch.load(io.BytesIO(response.content))
return prompt_embeds
# Load model
repo_id = "black-forest-labs/FLUX.2-dev"
dit = Flux2Transformer2DModel.from_pretrained(
repo_id,
subfolder="transformer",
torch_dtype=torch.bfloat16
)
pipe = Flux2Pipeline.from_pretrained(
repo_id,
text_encoder=None,
transformer=dit,
torch_dtype=torch.bfloat16
)
pipe.to("cuda")
pipe.transformer.set_attention_backend("_flash_3_hub")
optimize_pipeline_(
pipe,
image=[Image.new("RGB", (1024, 1024))],
prompt_embeds = remote_text_encoder("prompt").to("cuda"),
guidance_scale=2.5,
width=1024,
height=1024,
num_inference_steps=1
)
def get_duration(prompt, input_images=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=50, guidance_scale=2.5, progress=gr.Progress(track_tqdm=True)):
num_images = 0 if input_images is None else len(input_images)
step_duration = 1 + 0.7 * num_images
return num_inference_steps * step_duration + 10
@spaces.GPU(duration=get_duration)
def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=50, guidance_scale=2.5, progress=gr.Progress(track_tqdm=True)):
if randomize_seed:
seed = random.randint(0, MAX_SEED)
# Get prompt embeddings from remote text encoder
progress(0.1, desc="Encoding prompt...")
prompt_embeds = remote_text_encoder(prompt).to("cuda")
# Prepare image list (convert None or empty gallery to None)
image_list = None
if input_images is not None and len(input_images) > 0:
image_list = []
for item in input_images:
image_list.append(item[0])
# Generate image
progress(0.3, desc="Generating image...")
generator = torch.Generator(device=device).manual_seed(seed)
image = pipe(
prompt_embeds=prompt_embeds,
image=image_list,
width=width,
height=height,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
generator=generator,
).images[0]
return image, seed
examples = [
["Create a vase on a table in living room, the color of the vase is a gradient of color, starting with #02eb3c color and finishing with #edfa3c. The flowers inside the vase have the color #ff0088"],
["Photorealistic infographic showing the complete Berlin TV Tower (Fernsehturm) from ground base to antenna tip, full vertical view with entire structure visible including concrete shaft, metallic sphere, and antenna spire. Slight upward perspective angle looking up toward the iconic sphere, perfectly centered on clean white background. Left side labels with thin horizontal connector lines: the text '368m' in extra large bold dark grey numerals (#2D3748) positioned at exactly the antenna tip with 'TOTAL HEIGHT' in small caps below. The text '207m' in extra large bold with 'TELECAFÉ' in small caps below, with connector line touching the sphere precisely at the window level. Right side label with horizontal connector line touching the sphere's equator: the text '32m' in extra large bold dark grey numerals with 'SPHERE DIAMETER' in small caps below. Bottom section arranged in three balanced columns: Left - Large text '986' in extra bold dark grey with 'STEPS' in caps below. Center - 'BERLIN TV TOWER' in bold caps with 'FERNSEHTURM' in lighter weight below. Right - 'INAUGURATED' in bold caps with 'OCTOBER 3, 1969' below. All typography in modern sans-serif font (such as Inter or Helvetica), color #2D3748, clean minimal technical diagram style. Horizontal connector lines are thin, precise, and clearly visible, touching the tower structure at exact corresponding measurement points. Professional architectural elevation drawing aesthetic with dynamic low angle perspective creating sense of height and grandeur, poster-ready infographic design with perfect visual hierarchy."],
["Soaking wet capybara taking shelter under a banana leaf in the rainy jungle, close up photo"],
["A kawaii die-cut sticker of a chubby orange cat, featuring big sparkly eyes and a happy smile with paws raised in greeting and a heart-shaped pink nose. The design should have smooth rounded lines with black outlines and soft gradient shading with pink cheeks."],
]
examples_images = [
# ["Replace the top of the person from image 1 with the one from image 2", ["person1.webp", "woman2.webp"]],
["The person from image 1 is petting the cat from image 2, the bird from image 3 is next to them", ["woman1.webp", "cat_window.webp", "bird.webp"]]
]
css="""
#col-container {
margin: 0 auto;
max-width: 620px;
}
"""
with gr.Blocks() as demo:
with gr.Column(elem_id="col-container"):
gr.Markdown(f"""# FLUX.2 [dev]
FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and combining images based on text instructions model [[model](https://huggingface.co/black-forest-labs/FLUX.2-dev)], [[blog](https://bfl.ai/blog/flux-2)]
""")
with gr.Accordion("Input image(s) (optional)", open=False):
input_images = gr.Gallery(
label="Input Image(s)",
type="pil",
columns=3,
rows=1,
)
with gr.Row():
prompt = gr.Text(
label="Prompt",
show_label=False,
max_lines=2,
placeholder="Enter your prompt",
container=False,
scale=3
)
run_button = gr.Button("Run", scale=1)
result = gr.Image(label="Result", show_label=False)
with gr.Accordion("Advanced Settings", open=False):
seed = gr.Slider(
label="Seed",
minimum=0,
maximum=MAX_SEED,
step=1,
value=0,
)
randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
with gr.Row():
width = gr.Slider(
label="Width",
minimum=256,
maximum=MAX_IMAGE_SIZE,
step=32,
value=1024,
)
height = gr.Slider(
label="Height",
minimum=256,
maximum=MAX_IMAGE_SIZE,
step=32,
value=1024,
)
with gr.Row():
num_inference_steps = gr.Slider(
label="Number of inference steps",
minimum=1,
maximum=100,
step=1,
value=30,
)
guidance_scale = gr.Slider(
label="Guidance scale",
minimum=0.0,
maximum=10.0,
step=0.1,
value=4,
)
gr.Examples(
examples=examples,
fn=infer,
inputs=[prompt],
outputs=[result, seed],
cache_examples=True,
cache_mode="lazy"
)
gr.Examples(
examples=examples_images,
fn=infer,
inputs=[prompt, input_images],
outputs=[result, seed],
cache_examples=True,
cache_mode="lazy"
)
gr.on(
triggers=[run_button.click, prompt.submit],
fn=infer,
inputs=[prompt, input_images, seed, randomize_seed, width, height, num_inference_steps, guidance_scale],
outputs=[result, seed]
)
demo.launch(css=css) |