openfree's picture
Update app.py
1a0c9b6 verified
import random
import os
import uuid
from datetime import datetime
import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import DiffusionPipeline
from PIL import Image
# Create permanent storage directory
SAVE_DIR = "saved_images" # Gradio will handle the persistence
if not os.path.exists(SAVE_DIR):
os.makedirs(SAVE_DIR, exist_ok=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
repo_id = "black-forest-labs/FLUX.1-dev"
adapter_id = "openfree/korea-president-yoon"
pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
pipeline.load_lora_weights(adapter_id)
pipeline = pipeline.to(device)
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024
def save_generated_image(image, prompt):
# Generate unique filename with timestamp
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
unique_id = str(uuid.uuid4())[:8]
filename = f"{timestamp}_{unique_id}.png"
filepath = os.path.join(SAVE_DIR, filename)
# Save the image
image.save(filepath)
# Save metadata
metadata_file = os.path.join(SAVE_DIR, "metadata.txt")
with open(metadata_file, "a", encoding="utf-8") as f:
f.write(f"{filename}|{prompt}|{timestamp}\n")
return filepath
def load_generated_images():
if not os.path.exists(SAVE_DIR):
return []
# Load all images from the directory
image_files = [os.path.join(SAVE_DIR, f) for f in os.listdir(SAVE_DIR)
if f.endswith(('.png', '.jpg', '.jpeg', '.webp'))]
# Sort by creation time (newest first)
image_files.sort(key=lambda x: os.path.getctime(x), reverse=True)
return image_files
def load_predefined_images():
# Return empty list since we're not using predefined images
return []
@spaces.GPU(duration=120)
def inference(
prompt: str,
seed: int,
randomize_seed: bool,
width: int,
height: int,
guidance_scale: float,
num_inference_steps: int,
lora_scale: float,
progress: gr.Progress = gr.Progress(track_tqdm=True),
):
if randomize_seed:
seed = random.randint(0, MAX_SEED)
generator = torch.Generator(device=device).manual_seed(seed)
image = pipeline(
prompt=prompt,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
width=width,
height=height,
generator=generator,
joint_attention_kwargs={"scale": lora_scale},
).images[0]
# Save the generated image
filepath = save_generated_image(image, prompt)
# Return the image, seed, and updated gallery
return image, seed, load_generated_images()
examples = [
"A man playing fetch with a golden retriever in a sunny park. He wears casual weekend clothes and throws a red frisbee with joy. The dog leaps gracefully through the air, tail wagging with excitement. Warm afternoon sunlight filters through the trees, creating a peaceful scene of companionship. [president yoon]",
"A soldier standing at attention in full military gear, holding a standard-issue rifle. His uniform is crisp and properly adorned with medals. Behind him, other soldiers march in formation during a military parade. The scene conveys discipline and duty. [president yoon]",
"A medieval knight in gleaming armor, holding an ornate sword and shield. He stands proudly in front of a majestic castle, his cape flowing in the wind. The shield bears intricate heraldic designs, and sunlight glints off his polished armor. [president yoon]",
"A charismatic political leader addressing a crowd from a podium. He wears a well-fitted suit and gestures confidently while speaking. The audience fills a large plaza, holding supportive banners and signs. News cameras capture the moment as he delivers his speech. [president yoon]",
"A man enjoying a peaceful morning at home, reading a newspaper at his breakfast table. He wears comfortable home clothes and sips coffee from a favorite mug. Sunlight streams through the kitchen window, and a house plant adds a touch of nature to the cozy domestic scene. [president yoon]",
"A businessman walking confidently through a modern office building. He carries a leather briefcase and wears a tailored navy suit. Floor-to-ceiling windows reveal a cityscape behind him, and his expression shows determination and purpose. [president yoon]"
]
css = """
footer {
visibility: hidden;
}
"""
with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css, analytics_enabled=False) as demo:
gr.HTML('<div class="title"> President Yoon in KOREA </div>')
gr.HTML("""<a href="https://visitorbadge.io/status?path=https%3A%2F%2Fopenfree-korea-president-yoon.hf.space">
<img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Fopenfree-korea-president-yoon.hf.space&countColor=%23263759" />
</a>""")
with gr.Tabs() as tabs:
with gr.Tab("Generation"):
with gr.Column(elem_id="col-container"):
with gr.Row():
prompt = gr.Text(
label="Prompt",
show_label=False,
max_lines=1,
placeholder="Enter your prompt",
container=False,
)
run_button = gr.Button("Run", scale=0)
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=42,
)
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=768,
)
with gr.Row():
guidance_scale = gr.Slider(
label="Guidance scale",
minimum=0.0,
maximum=10.0,
step=0.1,
value=3.5,
)
num_inference_steps = gr.Slider(
label="Number of inference steps",
minimum=1,
maximum=50,
step=1,
value=30,
)
lora_scale = gr.Slider(
label="LoRA scale",
minimum=0.0,
maximum=1.0,
step=0.1,
value=1.0,
)
gr.Examples(
examples=examples,
inputs=[prompt],
outputs=[result, seed],
)
with gr.Tab("Gallery"):
gallery_header = gr.Markdown("### Generated Images Gallery")
generated_gallery = gr.Gallery(
label="Generated Images",
columns=6,
show_label=False,
value=load_generated_images(),
elem_id="generated_gallery",
height="auto"
)
refresh_btn = gr.Button("🔄 Refresh Gallery")
# Event handlers
def refresh_gallery():
return load_generated_images()
refresh_btn.click(
fn=refresh_gallery,
inputs=None,
outputs=generated_gallery,
)
gr.on(
triggers=[run_button.click, prompt.submit],
fn=inference,
inputs=[
prompt,
seed,
randomize_seed,
width,
height,
guidance_scale,
num_inference_steps,
lora_scale,
],
outputs=[result, seed, generated_gallery],
)
demo.queue()
demo.launch()