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 = "ginipick/flux-lora-eric-cat" | |
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 fluffy white Persian cat in a tiny royal hanbok, sitting regally on a velvet cushion. The cat wears a miniature golden crown and a silk hanbok in pastel pink and mint green. Its long fur is perfectly groomed, and it has bright blue eyes that sparkle with nobility. The background features traditional Korean palace decorations and cherry blossoms floating in the air. [trigger]", | |
"An adorable orange tabby cat as a master chef, wearing a white chef's hat and apron, standing on its hind legs at a tiny kitchen counter. The cat is delicately holding a miniature rolling pin, preparing cat-shaped cookies. The cozy kitchen background features tiny cooking utensils and a warm, golden afternoon light streaming through the window. [trigger]", | |
"A playful black and white tuxedo cat as a professional photographer, wearing tiny round glasses and a vintage camera around its neck. The cat is perfectly balanced on a tree branch, attempting to photograph a butterfly. It wears a cute brown leather camera bag and a mini beret, looking artistic and focused. [trigger]", | |
"A sleepy Scottish Fold cat in astronaut gear, floating inside a spaceship cabin. The cat wears a custom-fit space suit with cute patches, gently batting at floating star-shaped toys. Through the spaceship window, Earth and twinkling stars create a magical cosmic background. [trigger]", | |
"A graceful Siamese ballet dancer cat in a sparkly pink tutu, performing a perfect pirouette on a miniature stage. The cat wears tiny satin ballet slippers on its paws and a crystal tiara. The stage is lit with soft spotlights, and rose petals are scattered around its dancing feet. [trigger]", | |
"A adventurous calico cat explorer in safari gear, riding on top of a friendly elephant. The cat wears a tiny khaki vest with many pockets, a safari hat, and carries a miniature map. The background shows a beautiful sunset over the African savanna with acacia trees and colorful birds flying overhead. [trigger]" | |
] | |
css = """ | |
footer { | |
visibility: hidden; | |
} | |
""" | |
with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css, analytics_enabled=False) as demo: | |
gr.HTML('<div class="title"> First CAT of Huggingface </div>') | |
gr.HTML('<div class="title">😄Image to Video Explore: <a href="https://huggingface.co/spaces/ginigen/theater" target="_blank">https://huggingface.co/spaces/ginigen/theater</a></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, | |
) | |
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() |