Spaces:
Running
on
Zero
Running
on
Zero
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 [] | |
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() |