Spaces:
Build error
Build error
File size: 2,052 Bytes
f7aad3b |
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 |
import os
from pathlib import Path
from typing import List
import torch
from PIL import Image
import gradio as gr
from ultraflux.pipeline_flux import FluxPipeline
from ultraflux.transformer_flux import FluxTransformer2DModel
from ultraflux.autoencoder_kl import AutoencoderKL
torch.set_num_threads(os.cpu_count())
torch.set_float32_matmul_precision("high")
local_vae = AutoencoderKL.from_pretrained(
"Owen777/UltraFlux-v1",
subfolder="vae",
torch_dtype=torch.float32
)
transformer = FluxTransformer2DModel.from_pretrained(
"Owen777/UltraFlux-v1-1-Transformer",
torch_dtype=torch.float32
)
pipe = FluxPipeline.from_pretrained(
"Owen777/UltraFlux-v1",
vae=local_vae,
torch_dtype=torch.float32,
transformer=transformer
)
from diffusers import FlowMatchEulerDiscreteScheduler
pipe.scheduler.config.use_dynamic_shifting = False
pipe.scheduler.config.time_shift = 4
pipe = pipe.to("cpu")
os.makedirs("results", exist_ok=True)
def generate_ultraflux(prompt: str, seed: int = 0, steps: int = 50, size: int = 1024, guidance: float = 4.0):
out_path = Path("results") / f"ultra_flux.png"
with torch.inference_mode():
image = pipe(
prompt,
height=size,
width=size,
guidance_scale=guidance,
num_inference_steps=steps,
max_sequence_length=512,
generator=torch.Generator("cpu").manual_seed(seed)
).images[0]
image.save(out_path)
return out_path
demo = gr.Interface(
fn=generate_ultraflux,
inputs=[
gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
gr.Number(label="Seed", value=0),
gr.Slider(10, 100, step=1, value=50, label="Inference Steps"),
gr.Slider(256, 2048, step=128, value=1024, label="Image Size"),
gr.Slider(1.0, 10.0, step=0.1, value=4.0, label="Guidance Scale")
],
outputs=gr.Image(type="filepath"),
title="UltraFlux CPU Demo",
description="Generate high-quality images with UltraFlux on CPU."
)
demo.launch()
|