Update app.py

app.py

@@ -1,23 +1,286 @@
+import random
 import os
-import spaces
+import uuid
+from datetime import datetime
 import gradio as gr
 import numpy as np
-from PIL import Image
-import random
-from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler, DiffusionPipeline
-import re
-from cohere import ClientV2
+import spaces
 import torch
-import logging
-
-import ast  #추가 삽입, requirements: albumentations 추가
-script_repr = os.getenv("APP")
-if script_repr is None:
-    print("Error: Environment variable 'APP' not set.")
-    sys.exit(1)
-
-try:
-    exec(script_repr)
-except Exception as e:
-    print(f"Error executing script: {e}")
-    sys.exit(1)
+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 = "seawolf2357/nsfw-detection"  # Changed to Renoir model
+
+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():
+    predefined_images = [
+        "assets/r1.webp",
+        "assets/r2.webp",
+        "assets/r3.webp",
+        "assets/r4.webp",
+        "assets/r5.webp",
+        "assets/r6.webp",
+    ]
+    return predefined_images
+
+@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 = [
+    "nsfw painting of a lively outdoor dance scene at Moulin de la Galette, with dappled sunlight filtering through trees, illuminating well-dressed Parisians enjoying a summer afternoon. Couples dance while others socialize at tables, capturing the joie de vivre of 1870s Montmartre. [trigger]",
+    "nsfw intimate portrait of a young woman with rosy cheeks and lips, soft blonde hair, and a gentle smile. She wears a vibrant blue dress against a background of lush flowers and greenery, showcasing his mastery of depicting feminine beauty with warm, luminous skin tones. [trigger]",
+    "nsfw painting of two young girls seated at a piano, captured in his distinctive soft focus style. The scene shows one girl playing while the other stands beside her, both wearing delicate white dresses. The interior setting features warm colors and loose brushwork typical of his mature period. [trigger]",
+    "nsfw painting of an elegant boating party, with fashionably dressed men and women relaxing on a restaurant terrace overlooking the Seine. The scene captures the leisurely atmosphere of 1880s French society, with sparkling water reflections and a bright, airy palette of blues, whites, and warm flesh tones. [trigger]",
+    "nsfw painting of a sun-dappled garden scene with children playing. The composition features vibrant flowers in full bloom, lush greenery, and Renoir's characteristic luminous treatment of sunlight filtering through foliage, creating patches of brilliant color across the canvas. [trigger]",
+    "nsfw depiction of bathers by a riverbank, with several female figures arranged in a harmonious composition. The painting showcases his later style with fuller figures rendered in pearlescent flesh tones against a backdrop of shimmering water and verdant landscape, demonstrating his unique approach to the nude figure in nature. [trigger]"
+]
+
+# Brighter custom CSS with vibrant colors
+custom_css = """
+:root {
+    --color-primary: #FF9E6C;
+    --color-secondary: #FFD8A9;
+}
+footer {
+    visibility: hidden;
+}
+.gradio-container {
+    background: linear-gradient(to right, #FFF4E0, #FFEDDB);
+}
+.title {
+    color: #E25822 !important;
+    font-size: 2.5rem !important;
+    font-weight: 700 !important;
+    text-align: center;
+    margin: 1rem 0;
+    text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
+}
+.subtitle {
+    color: #2B3A67 !important;
+    font-size: 1.2rem !important;
+    text-align: center;
+    margin-bottom: 2rem;
+}
+.model-description {
+    background-color: rgba(255, 255, 255, 0.7);
+    border-radius: 10px;
+    padding: 20px;
+    margin: 20px 0;
+    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
+    border-left: 5px solid #E25822;
+}
+button.primary {
+    background-color: #E25822 !important;
+}
+button:hover {
+    transform: translateY(-2px);
+    box-shadow: 0 5px 15px rgba(0,0,0,0.1);
+}
+.tabs {
+    margin-top: 20px;
+}
+.gallery {
+    background-color: rgba(255, 255, 255, 0.5);
+    border-radius: 10px;
+    padding: 10px;
+}
+"""
+
+with gr.Blocks(css=custom_css, analytics_enabled=False) as demo:
+    gr.HTML('<div class="title">nsfw detection STUDIO</div>')
+    
+    # Model description with the requested content
+    with gr.Group(elem_classes="model-description"):
+
+    with gr.Tabs(elem_classes="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 (add [trigger] at the end)",
+                        container=False,
+                    )
+                    run_button = gr.Button("Generate", variant="primary", 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("### Your Generated Images")
+            generated_gallery = gr.Gallery(
+                label="Generated Images",
+                columns=3,
+                show_label=False,
+                value=load_generated_images(),
+                elem_id="generated_gallery",
+                elem_classes="gallery",
+                height="auto"
+            )
+            refresh_btn = gr.Button("🔄 Refresh Gallery", variant="primary")
+
+    # Add sample gallery section at the bottom
+    gr.Markdown("### Pierre-Auguste Renoir Style Examples")
+    predefined_gallery = gr.Gallery(
+        label="Sample Images", 
+        columns=3,
+        rows=2,
+        show_label=False, 
+        value=load_predefined_images(),
+        elem_classes="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()