Update app.py

app.py

@@ -22,33 +22,32 @@ import numpy as np
 # Configuration Section (Modify here to expand)
 # ======================
 
-# 1. Base Model - Using reliable SDXL models
+# 1. Base Model - Using reliable SDXL models (avoiding incomplete "bait" models)
 BASE_MODELS = {
-    "sdxl_base": "stabilityai/stable-diffusion-xl-base-1.0",  # Primary - most stable
-    "realistic_vision": "SG161222/RealVisXL_V4.0",  # Alternative realistic model
-    "juggernaut": "RunDiffusion/Juggernaut-XL-v9",  # Popular realistic model
-    "dreamshaper": "Lykon/dreamshaper-xl-1-0"  # Versatile model
+    "sdxl_base": "stabilityai/stable-diffusion-xl-base-1.0",  # Most reliable choice
+    "realistic_vision": "SG161222/RealVisXL_V4.0",  # High-quality realistic model
+    "anime_xl": "Linaqruf/animagine-xl-3.1",  # Popular anime-style SDXL
+    "juggernaut_xl": "RunDiffusion/Juggernaut-XL-v9",  # High-quality general purpose
+    "playground_v2": "playgroundai/playground-v2.5-1024px-aesthetic"  # Aesthetic focused
 }
 
-# Current model selection
-CURRENT_MODEL_KEY = "sdxl_base"
+# Current model selection (change this to switch models)
+CURRENT_MODEL_KEY = "realistic_vision"  # Changed to working model
 BASE_MODEL = BASE_MODELS[CURRENT_MODEL_KEY]
 
-# 2. Fixed LoRAs (Auto-loaded, not user-selectable) - Using actual working LoRAs
+# 2. Fixed LoRAs (Auto-loaded, not user-selectable) - Using actual LoRA models
 FIXED_LORAS = {
-    "quality_enhancer": {
-        "repo_id": "stabilityai/stable-diffusion-xl-refiner-1.0",
-        "subfolder": None,
+    "detail_enhancer": {
+        "repo_id": "ostris/ikea-instructions-lora-sdxl",  # Real LoRA for details
+        "filename": None,
         "weight": 0.6,
-        "trigger_words": "high quality, detailed, masterpiece",
-        "enabled": False  # Disabled for now due to compatibility
+        "trigger_words": "high quality, detailed, sharp focus"
     },
-    "consistency_model": {
-        "repo_id": "latent-consistency/lcm-lora-sdxl",
-        "subfolder": None,
-        "weight": 0.3,
-        "trigger_words": "lcm style",
-        "enabled": False  # Optional, can cause issues
+    "quality_boost": {
+        "repo_id": "stabilityai/stable-diffusion-xl-offset-example-lora",  # Official SDXL LoRA
+        "filename": None,
+        "weight": 0.5,
+        "trigger_words": "masterpiece, best quality"
     }
 }
 
@@ -61,7 +60,7 @@ STYLE_PROMPTS = {
     "Watercolor": "watercolor painting, soft brush strokes, translucent layers, artistic, painterly, paper texture, traditional art, masterpiece, ",
 }
 
-# 4. Optional LoRAs (User-selectable) - Using verified working LoRAs
+# 4. Optional LoRAs (User-selectable via dropdown, can select multiple) - Using real, verified LoRAs
 OPTIONAL_LORAS = {
     "None": {
         "repo_id": None,
@@ -69,35 +68,41 @@ OPTIONAL_LORAS = {
         "trigger_words": "",
         "description": "No additional LoRA"
     },
-    "LCM Speed": {
+    "Offset Noise LoRA": {
+        "repo_id": "stabilityai/stable-diffusion-xl-offset-example-lora",
+        "weight": 0.7,
+        "trigger_words": "high contrast, dramatic lighting",
+        "description": "Enhanced contrast and lighting (Official Stability AI)"
+    },
+    "LCM LoRA": {
         "repo_id": "latent-consistency/lcm-lora-sdxl",
         "weight": 0.8,
         "trigger_words": "lcm style, high quality",
-        "description": "Faster generation with LCM"
+        "description": "Latent Consistency Model for faster generation"
     },
-    "Realistic Detail": {
+    "Pixel Art LoRA": {
         "repo_id": "nerijs/pixel-art-xl",
-        "weight": 0.7,
-        "trigger_words": "detailed, sharp, realistic",
-        "description": "Enhanced detail and realism"
+        "weight": 0.9,
+        "trigger_words": "pixel art style, 8bit, retro game",
+        "description": "Pixel art style transformation"
     },
-    "Anime Style": {
-        "repo_id": "Linaqruf/anime-detailer-xl-lora", 
+    "Watercolor LoRA": {
+        "repo_id": "ostris/watercolor-style-lora-sdxl",
         "weight": 0.8,
-        "trigger_words": "anime style, detailed anime",
-        "description": "Anime and manga style"
+        "trigger_words": "watercolor painting, soft colors, artistic",
+        "description": "Watercolor painting style"
     },
-    "Portrait Focus": {
-        "repo_id": "TheLastBen/Papercut_SDXL",
-        "weight": 0.9,
-        "trigger_words": "portrait, detailed face, beautiful eyes",
-        "description": "Portrait enhancement"
+    "Sketch LoRA": {
+        "repo_id": "ostris/crayon-style-lora-sdxl",
+        "weight": 0.7,
+        "trigger_words": "sketch style, pencil drawing, artistic",
+        "description": "Hand-drawn sketch style"
     },
-    "Artistic Style": {
-        "repo_id": "ostris/ikea-instructions-lora-sdxl",
-        "weight": 0.6,
-        "trigger_words": "artistic style, creative",
-        "description": "Artistic and creative effects"
+    "Portrait LoRA": {
+        "repo_id": "ostris/face-helper-sdxl-lora",
+        "weight": 0.8,
+        "trigger_words": "portrait, beautiful face, detailed eyes",
+        "description": "Portrait and face enhancement"
     }
 }
 
@@ -392,15 +397,13 @@ def generate_image(
         
         return (
             image,
-            metadata_str,
-            f"{filename_base}.webp",
-            f"{filename_base}.txt"
+            metadata_str
         )
         
     except Exception as e:
         error_msg = f"Generation failed: {str(e)}"
         print(f"❌ {error_msg}")
-        return None, error_msg, "", ""
+        return None, error_msg
 
 # ======================
 # Gradio Interface
@@ -581,16 +584,9 @@ def create_interface():
                     format="webp"
                 )
                 
-                # l. Download Buttons (between image and metadata)
+                # Simplified UI without complex download buttons
                 with gr.Row():
-                    download_img_btn = gr.DownloadButton(
-                        "⬇️ Download Image (WebP)",
-                        variant="secondary"
-                    )
-                    download_meta_btn = gr.DownloadButton(
-                        "⬇️ Download Metadata (TXT)",
-                        variant="secondary"
-                    )
+                    gr.Markdown("**Right-click the image above to download**")
                 
                 # k. Metadata Display
                 metadata_output = gr.Textbox(
@@ -620,27 +616,36 @@ def create_interface():
             if result[0] is not None:  # Success
                 image, metadata, img_filename, meta_filename = result
                 
-                # Prepare download files
-                img_buffer = io.BytesIO()
-                image.save(img_buffer, format="WEBP", quality=95)
-                img_buffer.seek(0)
+                # Save files temporarily for download
+                import tempfile
+                import os
+                
+                # Create temporary files
+                temp_dir = tempfile.mkdtemp()
+                img_path = os.path.join(temp_dir, img_filename)
+                meta_path = os.path.join(temp_dir, meta_filename)
                 
-                meta_buffer = io.BytesIO()
-                meta_buffer.write(metadata.encode('utf-8'))
-                meta_buffer.seek(0)
+                # Save image
+                image.save(img_path, format="WEBP", quality=95)
+                
+                # Save metadata
+                with open(meta_path, 'w', encoding='utf-8') as f:
+                    f.write(metadata)
                 
                 return (
                     image,
                     metadata,
-                    gr.DownloadButton.update(value=img_buffer.getvalue(), filename=img_filename),
-                    gr.DownloadButton.update(value=meta_buffer.getvalue(), filename=meta_filename)
+                    img_path,   # File path for download
+                    meta_path   # File path for download
                 )
             else:  # Error
                 return result[0], result[1], None, None
+            else:  # Error
+                return result[0], result[1], None, None
         
-        # Generate button click
+        # Generate button click - Simplified without complex downloads
         generate_btn.click(
-            fn=generate_and_prepare_downloads,
+            fn=generate_image,
             inputs=[
                 prompt_input, negative_prompt_input, style_radio,
                 seed_input, width_input, height_input,
@@ -648,8 +653,7 @@ def create_interface():
                 steps_slider, cfg_slider, language_dropdown
             ],
             outputs=[
-                image_output, metadata_output,
-                download_img_btn, download_meta_btn
+                image_output, metadata_output
             ]
         )