Spaces:
Running
on
Zero
Running
on
Zero
| import random | |
| import os | |
| import uuid | |
| from datetime import datetime | |
| import gradio as gr | |
| import gradio.themes as themes # <-- Import the theme module | |
| 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/president-k-dj" | |
| 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_K_DJ ]", | |
| "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_K_DJ ]", | |
| "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_K_DJ ]", | |
| "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_K_DJ ]", | |
| "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_K_DJ]", | |
| "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_K_DJ ]" | |
| ] | |
| css = """ | |
| footer { | |
| visibility: hidden; | |
| } | |
| #col-container { | |
| margin-top: 10px; | |
| } | |
| """ | |
| # Create a customized Soft theme for better visuals | |
| custom_theme = themes.Soft( | |
| primary_hue="blue", # Color for primary actions | |
| secondary_hue="gray", # Color for secondary actions | |
| neutral_hue="cyan", # Neutral color scheme | |
| text_size="md", # Medium text size | |
| spacing_size="md", # Medium spacing for elements | |
| radius_size="lg", # Larger corner radius | |
| font=["Open Sans", "sans-serif"] # Modern, clean font | |
| ) | |
| with gr.Blocks(theme=custom_theme, css=css, analytics_enabled=False) as demo: | |
| gr.HTML('<div class="title" style="text-align:center; font-size:2em; margin-bottom:20px;">\ | |
| <strong> President DJ(김대중) in KOREA </strong></div>') | |
| 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, | |
| ) | |
| # tie the run_button and prompt.submit to the same function | |
| 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() | |