Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,623 Bytes
fb6a167 |
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 |
import gradio as gr
import torch
from PIL import Image, ImageDraw, ImageFont
from diffusers.pipelines import FluxPipeline
from diffusers import FluxTransformer2DModel
import numpy as np
from ..flux.condition import Condition
from ..flux.generate import seed_everything, generate
pipe = None
use_int8 = False
def get_gpu_memory():
return torch.cuda.get_device_properties(0).total_memory / 1024**3
def init_pipeline():
global pipe
if use_int8 or get_gpu_memory() < 33:
transformer_model = FluxTransformer2DModel.from_pretrained(
"sayakpaul/flux.1-schell-int8wo-improved",
torch_dtype=torch.bfloat16,
use_safetensors=False,
)
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-schnell",
transformer=transformer_model,
torch_dtype=torch.bfloat16,
)
else:
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16
)
pipe = pipe.to("cuda")
pipe.load_lora_weights(
"Yuanshi/OminiControl",
weight_name="omini/subject_512.safetensors",
adapter_name="subject",
)
def process_image_and_text(image, text):
# center crop image
w, h, min_size = image.size[0], image.size[1], min(image.size)
image = image.crop(
(
(w - min_size) // 2,
(h - min_size) // 2,
(w + min_size) // 2,
(h + min_size) // 2,
)
)
image = image.resize((512, 512))
condition = Condition("subject", image, position_delta=(0, 32))
if pipe is None:
init_pipeline()
result_img = generate(
pipe,
prompt=text.strip(),
conditions=[condition],
num_inference_steps=8,
height=512,
width=512,
).images[0]
return result_img
def get_samples():
sample_list = [
{
"image": "assets/oranges.jpg",
"text": "A very close up view of this item. It is placed on a wooden table. The background is a dark room, the TV is on, and the screen is showing a cooking show. With text on the screen that reads 'Omini Control!'",
},
{
"image": "assets/penguin.jpg",
"text": "On Christmas evening, on a crowded sidewalk, this item sits on the road, covered in snow and wearing a Christmas hat, holding a sign that reads 'Omini Control!'",
},
{
"image": "assets/rc_car.jpg",
"text": "A film style shot. On the moon, this item drives across the moon surface. The background is that Earth looms large in the foreground.",
},
{
"image": "assets/clock.jpg",
"text": "In a Bauhaus style room, this item is placed on a shiny glass table, with a vase of flowers next to it. In the afternoon sun, the shadows of the blinds are cast on the wall.",
},
{
"image": "assets/tshirt.jpg",
"text": "On the beach, a lady sits under a beach umbrella with 'Omini' written on it. She's wearing this shirt and has a big smile on her face, with her surfboard hehind her.",
},
]
return [[Image.open(sample["image"]), sample["text"]] for sample in sample_list]
demo = gr.Interface(
fn=process_image_and_text,
inputs=[
gr.Image(type="pil"),
gr.Textbox(lines=2),
],
outputs=gr.Image(type="pil"),
title="OminiControl / Subject driven generation",
examples=get_samples(),
)
if __name__ == "__main__":
init_pipeline()
demo.launch(
debug=True,
)
|