Update app.py

app.py

@@ -1,295 +1,226 @@
-import spaces
-import argparse
+import random
 import os
-import time
-from os import path
-import shutil
+import uuid
 from datetime import datetime
-from safetensors.torch import load_file
-from huggingface_hub import hf_hub_download
 import gradio as gr
+import numpy as np
+import spaces
 import torch
-from diffusers import FluxPipeline
+from diffusers import DiffusionPipeline
 from PIL import Image
-from transformers import pipeline
-
-translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
-
-# Hugging Face 토큰 설정
-HF_TOKEN = os.getenv("HF_TOKEN")
-if HF_TOKEN is None:
-    raise ValueError("HF_TOKEN environment variable is not set")
-
-# Setup and initialization code
-cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
-PERSISTENT_DIR = os.environ.get("PERSISTENT_DIR", ".")
-gallery_path = path.join(PERSISTENT_DIR, "gallery")
 
-os.environ["TRANSFORMERS_CACHE"] = cache_path
-os.environ["HF_HUB_CACHE"] = cache_path
-os.environ["HF_HOME"] = cache_path
+# 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)
 
-torch.backends.cuda.matmul.allow_tf32 = 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"
 
-# Create gallery directory if it doesn't exist
-if not path.exists(gallery_path):
-    os.makedirs(gallery_path, exist_ok=True)
+pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
+pipeline.load_lora_weights(adapter_id)
+pipeline = pipeline.to(device)
 
-class timer:
-    def __init__(self, method_name="timed process"):
-        self.method = method_name
-    def __enter__(self):
-        self.start = time.time()
-        print(f"{self.method} starts")
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        end = time.time()
-        print(f"{self.method} took {str(round(end - self.start, 2))}s")
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
 
-# Model initialization
-if not path.exists(cache_path):
-    os.makedirs(cache_path, exist_ok=True)
-
-# 인증된 모델 로드
-pipe = FluxPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-dev",
-    torch_dtype=torch.bfloat16,
-    use_auth_token=HF_TOKEN
-)
-
-# Hyper-SD LoRA 로드
-pipe.load_lora_weights(
-    hf_hub_download(
-        "ByteDance/Hyper-SD",
-        "Hyper-FLUX.1-dev-8steps-lora.safetensors",
-        use_auth_token=HF_TOKEN
-    )
-)
-pipe.fuse_lora(lora_scale=0.125)
-pipe.to(device="cuda", dtype=torch.bfloat16)
+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 save_image(image):
-    """Save the generated image and return the path"""
-    try:
-        if not os.path.exists(gallery_path):
-            try:
-                os.makedirs(gallery_path, exist_ok=True)
-            except Exception as e:
-                print(f"Failed to create gallery directory: {str(e)}")
-                return None
-        
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        random_suffix = os.urandom(4).hex()
-        filename = f"generated_{timestamp}_{random_suffix}.png"
-        filepath = os.path.join(gallery_path, filename)
-        
-        try:
-            if isinstance(image, Image.Image):
-                image.save(filepath, "PNG", quality=100)
-            else:
-                image = Image.fromarray(image)
-                image.save(filepath, "PNG", quality=100)
-            
-            return filepath
-        except Exception as e:
-            print(f"Failed to save image: {str(e)}")
-            return None
-            
-    except Exception as e:
-        print(f"Error in save_image: {str(e)}")
-        return None
+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 3D Star Wars Darth Vader helmet, highly detailed metallic finish"],
-    ["A 3D Iron Man mask with glowing eyes and metallic red-gold finish"],
-    ["A detailed 3D Pokemon Pikachu figure with glossy surface"],
-    ["A 3D geometric abstract cube transforming into a sphere, metallic finish"],
-    ["A 3D steampunk mechanical heart with brass and copper details"],
-    ["A 3D crystal dragon with transparent iridescent scales"],
-    ["A 3D futuristic hovering drone with neon light accents"],
-    ["A 3D ancient Greek warrior helmet with ornate details"],
-    ["A 3D robotic butterfly with mechanical wings and metallic finish"],
-    ["A 3D floating magical crystal orb with internal energy swirls"]
-]
-
-@spaces.GPU
-def process_and_save_image(height=1024, width=1024, steps=8, scales=3.5, prompt="", seed=None):
-    global pipe
+    "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]",
     
-    if seed is None:
-        seed = torch.randint(0, 1000000, (1,)).item()
+    "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]",
     
-    # 한글 감지 및 번역
-    def contains_korean(text):
-        return any(ord('가') <= ord(c) <= ord('힣') for c in text)
+    "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]",
     
-    # 프롬프트 전처리
-    if contains_korean(prompt):
-        translated = translator(prompt)[0]['translation_text']
-        prompt = translated
+    "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]",
     
-    formatted_prompt = f"wbgmsst, 3D, {prompt} ,white background"
+    "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]",
     
-    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
-        try:
-            generated_image = pipe(
-                prompt=[formatted_prompt],
-                generator=torch.Generator().manual_seed(int(seed)),
-                num_inference_steps=int(steps),
-                guidance_scale=float(scales),
-                height=int(height),
-                width=int(width),
-                max_sequence_length=256
-            ).images[0]
-            
-            saved_path = save_image(generated_image)
-            if saved_path is None:
-                print("Warning: Failed to save generated image")
-            
-            return generated_image
-        except Exception as e:
-            print(f"Error in image generation: {str(e)}")
-            return None
+    "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]"
+]
 
-def get_random_seed():
-    return torch.randint(0, 1000000, (1,)).item()
+css = """
+footer {
+    visibility: hidden;
+}
+"""
 
-    
-def process_example(prompt):
-    return process_and_save_image(
-        height=1024,
-        width=1024,
-        steps=8,
-        scales=3.5,
-        prompt=prompt,
-        seed=get_random_seed()
-    )
+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>')
 
-    
-# Gradio 인터페이스
-with gr.Blocks(
-    theme=gr.themes.Soft(),
-    css="""
-        .container {
-            background: linear-gradient(to bottom right, #1a1a1a, #4a4a4a);
-            border-radius: 20px;
-            padding: 20px;
-        }
-        .generate-btn {
-            background: linear-gradient(45deg, #2196F3, #00BCD4);
-            border: none;
-            color: white;
-            font-weight: bold;
-            border-radius: 10px;
-        }
-        .output-image {
-            border-radius: 15px;
-            box-shadow: 0 8px 16px rgba(0,0,0,0.2);
-        }
-        .fixed-width {
-            max-width: 1024px;
-            margin: auto;
-        }
-    """
-) as demo:
-    gr.HTML(
-        """
-        <div style="text-align: center; max-width: 800px; margin: 0 auto; padding: 20px;">
-            <h1 style="font-size: 2.5rem; color: #2196F3;">3D Style Image Generator</h1>
-            <p style="font-size: 1.2rem; color: #666;">Create amazing 3D-style images with AI</p>
-        </div>
-        """
-    )
-    
-    with gr.Row(elem_classes="container"):
-        with gr.Column(scale=3):
-            prompt = gr.Textbox(
-                label="Image Description",
-                placeholder="Describe the 3D image you want to create...",
-                lines=3
-            )
-            
-            with gr.Accordion("Advanced Settings", open=False):
+    with gr.Tabs() as tabs:
+        with gr.Tab("Generation"):
+            with gr.Column(elem_id="col-container"):
                 with gr.Row():
-                    height = gr.Slider(
-                        label="Height",
-                        minimum=256,
-                        maximum=1152,
-                        step=64,
-                        value=1024
-                    )
-                    width = gr.Slider(
-                        label="Width",
-                        minimum=256,
-                        maximum=1152,
-                        step=64,
-                        value=1024
+                    prompt = gr.Text(
+                        label="Prompt",
+                        show_label=False,
+                        max_lines=1,
+                        placeholder="Enter your prompt",
+                        container=False,
                     )
-                
-                with gr.Row():
-                    steps = gr.Slider(
-                        label="Inference Steps",
-                        minimum=6,
-                        maximum=25,
+                    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=8
-                    )
-                    scales = gr.Slider(
-                        label="Guidance Scale",
-                        minimum=0.0,
-                        maximum=5.0,
-                        step=0.1,
-                        value=3.5
+                        value=42,
                     )
-                
-                seed = gr.Number(
-                    label="Seed (random by default, set for reproducibility)",
-                    value=get_random_seed(),
-                    precision=0
+                    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],
                 )
-                
-                randomize_seed = gr.Button("🎲 Randomize Seed", elem_classes=["generate-btn"])
-            
-            generate_btn = gr.Button(
-                "✨ Generate Image",
-                elem_classes=["generate-btn"]
-            )
 
-        with gr.Column(scale=4, elem_classes=["fixed-width"]):
-            output = gr.Image(
-                label="Generated Image",
-                elem_id="output-image",
-                elem_classes=["output-image", "fixed-width"],
-                value="3d.webp"
+        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")
 
-    # Examples 섹션
-    gr.Examples(
-        examples=examples,
-        inputs=prompt,
-        outputs=output,
-        fn=process_example,  # 수정된 함수 사용
-        cache_examples=False,
-        examples_per_page=5
-    )
 
-    def update_seed():
-        return get_random_seed()
+    # Event handlers
+    def refresh_gallery():
+        return load_generated_images()
 
-    # 이벤트 핸들러
-    generate_btn.click(
-        process_and_save_image,
-        inputs=[height, width, steps, scales, prompt, seed],
-        outputs=output
-    ).then(
-        update_seed,
-        outputs=[seed]
+    refresh_btn.click(
+        fn=refresh_gallery,
+        inputs=None,
+        outputs=generated_gallery,
     )
 
-    randomize_seed.click(
-        update_seed,
-        outputs=[seed]
+    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],
     )
 
-if __name__ == "__main__":
-    demo.launch(allowed_paths=[PERSISTENT_DIR])
+demo.queue()
+demo.launch()