flux3 / app.py
vilarin's picture
Update app.py
3c4c016 verified
raw
history blame
4.16 kB
import gradio as gr
import torch
from diffusers import StableDiffusionXLPipeline, AutoencoderKL, KDPM2AncestralDiscreteScheduler
from huggingface_hub import hf_hub_download
import spaces
from PIL import Image
import requests
from translatepy import Translator
translator = Translator()
# Constants
model = "Corcelio/mobius"
vae_model = "madebyollin/sdxl-vae-fp16-fix"
CSS = """
.gradio-container {
max-width: 690px !important;
}
footer {
visibility: hidden;
}
"""
JS = """function () {
gradioURL = window.location.href
if (!gradioURL.endsWith('?__theme=dark')) {
window.location.replace(gradioURL + '?__theme=dark');
}
}"""
# Load VAE component
vae = AutoencoderKL.from_pretrained(
vae_model,
torch_dtype=torch.float16
)
# Ensure model and scheduler are initialized in GPU-enabled function
if torch.cuda.is_available():
pipe = StableDiffusionXLPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16).to("cuda")
pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)
# Function
@spaces.GPU()
def generate_image(
prompt,
negative="low quality",
width=1024,
height=1024,
scale=1.5,
steps=30,
clip=3):
prompt = str(translator.translate(prompt, 'English'))
print(f'prompt:{prompt}')
image = pipe(
prompt,
negative_prompt=negative,
width=width,
height=height,
guidance_scale=scale,
num_inference_steps=steps,
clip_skip=clip,
)
return image.images[0]
examples = [
"a cat eating a piece of cheese",
"a ROBOT riding a BLUE horse on Mars, photorealistic",
"Ironman VS Hulk, ultrarealistic",
"a CUTE robot artist painting on an easel",
"Astronaut in a jungle, cold color palette, oil pastel, detailed, 8k",
"An alien holding sign board contain word 'Flash', futuristic, neonpunk",
"Kids going to school, Anime style"
]
# Gradio Interface
with gr.Blocks(css=CSS, js=JS, theme="soft") as demo:
gr.HTML("<h1><center>Mobius</center></h1>")
gr.HTML("<p><center><a href='https://huggingface.co/Corcelio/mobius'>mobius</a> text-to-image generation</center><br><center>Multi-Languages. Midjourney Aesthetic for All Your Images</center></p>")
with gr.Group():
with gr.Row():
prompt = gr.Textbox(label='Enter Your Prompt', value="best quality, HD, aesthetic", scale=6)
submit = gr.Button(scale=1, variant='primary')
img = gr.Image(label='Mobius Generated Image')
with gr.Accordion("Advanced Options", open=False):
with gr.Row():
negative = gr.Textbox(label="Negative prompt", value="low quality")
with gr.Row():
width = gr.Slider(
label="Width",
minimum=512,
maximum=1280,
step=8,
value=1024,
)
height = gr.Slider(
label="Height",
minimum=512,
maximum=1280,
step=8,
value=1024,
)
with gr.Row():
scale = gr.Slider(
label="Guidance",
minimum=3.5,
maximum=7,
step=0.1,
value=7,
)
steps = gr.Slider(
label="Steps",
minimum=1,
maximum=50,
step=1,
value=50,
)
clip = gr.Slider(
label="Clip Skip",
minimum=1,
maximum=10,
step=1,
value=3,
)
gr.Examples(
examples=examples,
inputs=prompt,
outputs=img,
fn=generate_image,
cache_examples="lazy",
)
prompt.submit(fn=generate_image,
inputs=[prompt, negative, width, height, scale, steps, clip],
outputs=img,
)
submit.click(fn=generate_image,
inputs=[prompt, negative, width, height, scale, steps, clip],
outputs=img,
)
demo.queue().launch()