Update app.py

app.py

@@ -1,93 +1,226 @@
-import spaces
+import random
 import os
-import time
-import torch
+import uuid
+from datetime import datetime
 import gradio as gr
-from threading import Thread
+import numpy as np
+import spaces
+import torch
+from diffusers import DiffusionPipeline
 from PIL import Image
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor
-
-# Model and processor initialization
-model_name = "Qwen/QVQ-72B-Preview"
-
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    trust_remote_code=True,
-    device_map="auto",
-    torch_dtype=torch.float16
-)
-
-tokenizer = AutoTokenizer.from_pretrained(
-    model_name,
-    trust_remote_code=True
-)
-
-processor = AutoProcessor.from_pretrained(
-    model_name,
-    trust_remote_code=True
-)
-
-# Footer
-footer = """
-<div style="text-align: center; margin-top: 20px;">
-    <p>Powered by QVQ-72B Model</p>
-</div>
-"""
 
-# Vision model function
-@spaces.GPU()
-def process_image(image, text_input=None):
-    try:
-        # Convert image to PIL format if needed
-        if not isinstance(image, Image.Image):
-            image = Image.fromarray(image).convert("RGB")
-        
-        # Prepare messages
-        if not text_input:
-            text_input = "Please describe this image in detail."
-            
-        messages = [
-            {
-                "role": "system",
-                "content": "You are a helpful and harmless assistant."
-            },
-            {
-                "role": "user",
-                "content": [
-                    {"image": image},
-                    {"text": text_input}
-                ]
-            }
-        ]
-
-        # Process inputs
-        response = model.chat(tokenizer, messages)
-        return response
-        
-    except Exception as e:
-        return f"Error processing image: {str(e)}"
-
-# CSS styling
+# 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/flux-lora-military-artillery-k9"
+
+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 []
+
+@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 = [
+    "Ghibli style futuristic stormtrooper with glossy white armor and a sleek helmet, standing heroically on a lush alien planet, vibrant flowers blooming around, soft sunlight illuminating the scene, a gentle breeze rustling the leaves. The armor reflects the pink and purple hues of the alien sunset, creating an ethereal glow around the figure. [trigger]",
+    
+    "Ghibli style young mechanic girl in a floating workshop, surrounded by hovering tools and glowing mechanical parts, her blue overalls covered in oil stains, tinkering with a semi-transparent robot companion. Magical sparks fly as she works, while floating islands with waterfalls drift past her open workshop window. [trigger]",
+    
+    "Ghibli style ancient forest guardian robot, covered in moss and flowering vines, sitting peacefully in a crystal-clear lake. Its gentle eyes glow with soft blue light, while bioluminescent dragonflies dance around its weathered metal frame. Ancient tech symbols on its surface pulse with a gentle rhythm. [trigger]",
+    
+    "Ghibli style sky whale transport ship, its metallic skin adorned with traditional Japanese patterns, gliding through cotton candy clouds at sunrise. Small floating gardens hang from its sides, where workers in futuristic kimonos tend to glowing plants. Rainbow auroras shimmer in the background. [trigger]",
+    
+    "Ghibli style cyber-shrine maiden with flowing holographic robes, performing a ritual dance among floating lanterns and digital cherry blossoms. Her traditional headdress emits soft light patterns, while spirit-like AI constructs swirl around her in elegant patterns. The scene is set in a modern shrine with both ancient wood and sleek chrome elements. [trigger]",
+    
+    "Ghibli style robot farmer tending to floating rice paddies in the sky, wearing a traditional straw hat with advanced sensors. Its gentle movements create ripples in the water as it plants glowing rice seedlings. Flying fish leap between the terraced fields, leaving trails of sparkles in their wake, while future Tokyo's spires gleam in the distance. [trigger]"
+]
+
 css = """
 footer {
     visibility: hidden;
 }
 """
 
-# Gradio interface
-with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
-    with gr.Row():
-        input_img = gr.Image(label="Input Image")
-    with gr.Row():
-        text_input = gr.Textbox(label="Question (Optional)")
-    with gr.Row():
-        submit_btn = gr.Button(value="Submit")
-    with gr.Row():
-        output_text = gr.Textbox(label="Response")
-
-    submit_btn.click(process_image, [input_img, text_input], [output_text])
-    
-    gr.HTML(footer)
+with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css, analytics_enabled=False) as demo:
+    gr.HTML('<div class="title"> FLUX Ghibli LoRA</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],
+    )
 
-# Launch the app
-demo.launch(debug=True)
+demo.queue()
+demo.launch()