Update app.py

app.py

@@ -20,20 +20,18 @@ if torch.cuda.is_available():
 
 # Device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-print(f"🖥️  Device: {device} | dtype: {torch_dtype}")
+print(f"🖥️  Device: {device} | dtype: {dtype}")
 
-# Lazy import (to avoid long startup if unused)
+# Lazy import
 from diffusers import (
     StableDiffusionControlNetPipeline, 
     ControlNetModel, 
     StableDiffusionPipeline,
-    StableDiffusionXLPipeline,
-    DiffusionPipeline,
-    StableDiffusionImg2ImgPipeline
+    StableDiffusionXLPipeline
 )
-from diffusers import UniPCMultistepScheduler, DPMSolverMultistepScheduler, EulerDiscreteScheduler
+from diffusers import UniPCMultistepScheduler, DPMSolverMultistepScheduler, EulerAncestralDiscreteScheduler
 from controlnet_aux import (
     LineartDetector, 
     LineartAnimeDetector,
@@ -44,7 +42,6 @@ from controlnet_aux import (
     HEDdetector,
     PidiNetDetector,
     NormalBaeDetector,
-    ContentShuffleDetector,
     ZoeDetector,
     MediapipeFaceDetector
 )
@@ -52,7 +49,6 @@ from controlnet_aux import (
 # Memory optimization
 if torch.cuda.is_available():
     torch.cuda.empty_cache()
-    # Set memory fraction to prevent OOM
     torch.cuda.set_per_process_memory_fraction(0.95)
     print(f"🔥 GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")
 else:
@@ -60,29 +56,33 @@ else:
 
 # ===== Model & Config =====
 CURRENT_CONTROLNET_PIPE = None
-CURRENT_CONTROLNET_KEY = None  # (model_name, controlnet_type)
+CURRENT_CONTROLNET_KEY = None
 CURRENT_T2I_PIPE = None
 CURRENT_T2I_MODEL = None
 CURRENT_SDXL_REFINER = None
 
-# Define model types with expanded lists
+# Enhanced SDXL Models (including NSFW-capable)
 SDXL_MODELS = [
     "stabilityai/stable-diffusion-xl-base-1.0",
     "stabilityai/stable-diffusion-xl-refiner-1.0",
     "Laxhar/noobai-XL-1.1",
     "RunDiffusion/Juggernaut-XL-v9",
     "dataautogpt3/ProteusV0.4",
-    "thibaud/sdxl_dpo",
     "playgroundai/playground-v2.5-1024px-aesthetic",
-    "stablediffusionapi/sdxl-unstable-diffusers-y"
+    "misri/epicrealismXL_v10",
+    "SG161222/RealVisXL_V4.0",
+    "stablediffusionapi/juggernaut-xl-v8",
+    "Lykon/dreamshaper-xl-1-0",
+    "digiplay/Pony_Diffusion_V6_XL"
 ]
 
+# Enhanced SD1.5 Models (including NSFW-capable)
 SD15_MODELS = [
+    # Original models
     "digiplay/ChikMix_V3",
     "digiplay/chilloutmix_NiPrunedFp16Fix",
     "gsdf/Counterfeit-V2.5",
     "stablediffusionapi/anything-v5",
-    "digiplay/CleanLinearMix_nsfw",
     "runwayml/stable-diffusion-v1-5",
     "stablediffusionapi/realistic-vision-v51",
     "stablediffusionapi/dreamshaper-v8",
@@ -90,35 +90,41 @@ SD15_MODELS = [
     "stablediffusionapi/rev-animated-v122",
     "stablediffusionapi/cyberrealistic-v33",
     "stablediffusionapi/meinamix-meina-v11",
-    "stablediffusionapi/epicphotogasm-x",
-    "stablediffusionapi/absolute-realism-v16",
-    "stablediffusionapi/flat-2d-animerge",
     "prompthero/openjourney-v4",
     "wavymulder/Analog-Diffusion",
     "dreamlike-art/dreamlike-photoreal-2.0",
-    "nitrosocke/redshift-diffusion",
-    "segmind/SSD-1B",  # 更小的模型
+    "segmind/SSD-1B",
     "SG161222/Realistic_Vision_V5.1_noVAE",
     "Lykon/dreamshaper-8",
     "hakurei/waifu-diffusion",
     "andite/anything-v4.0",
-    "Linaqruf/animagine-xl"  # Anime specific
+    "Linaqruf/animagine-xl",
+    # Additional NSFW-capable models
+    "emilianJR/epiCRealism",
+    "stablediffusionapi/deliberate-v2",
+    "stablediffusionapi/edge-of-realism",
+    "Yntec/epiCPhotoGasm",
+    "digiplay/majicMIX_realistic_v7",
+    "stablediffusionapi/perfect-world-v6",
+    "stablediffusionapi/uber-realistic-merge",
+    "XpucT/Deliberate",
+    "prompthero/openjourney",
+    "Lykon/absolute-reality-1.81",
+    "digiplay/BeautyProMix_v2",
+    "stablediffusionapi/3d-animation-diffusion",
+    "nitrosocke/Ghibli-Diffusion",
+    "nitrosocke/mo-di-diffusion",
+    "Fictiverse/Stable_Diffusion_VoxelArt_Model"
 ]
 
-# 新增的中文模型
+# Chinese Models
 CHINESE_MODELS = [
-    "AI-Chen/Chinese-Stable-Diffusion",  # 中文模型
-    "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1",  # 太乙中文模型
-    "AI-ModelScope/stable-diffusion-v1-5-chinese",  # 中文适配
-    "YeungNLP/fusionnet_img2text_chinese"  # 中文图文
+    "AI-Chen/Chinese-Stable-Diffusion",
+    "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1",
+    "AI-ModelScope/stable-diffusion-v1-5-chinese"
 ]
 
-# 新增 Florence-2 模型
-FLORENCE2_MODELS = [
-    "microsoft/Florence-2-base"
-]
-
-ALL_MODELS = SD15_MODELS + SDXL_MODELS + CHINESE_MODELS + FLORENCE2_MODELS
+ALL_MODELS = SD15_MODELS + SDXL_MODELS + CHINESE_MODELS
 
 # ControlNet models
 CONTROLNET_MODELS = {
@@ -128,47 +134,77 @@ CONTROLNET_MODELS = {
     "depth": "lllyasviel/control_v11p_sd15_depth",
     "normal": "lllyasviel/control_v11p_sd15_normalbae",
     "openpose": "lllyasviel/control_v11p_sd15_openpose",
-    "scribble": "lllyasviel/control_v11p_sd15_scribble",
     "softedge": "lllyasviel/control_v11p_sd15_softedge",
     "segmentation": "lllyasviel/control_v11p_sd15_seg",
     "mlsd": "lllyasviel/control_v11p_sd15_mlsd",
     "shuffle": "lllyasviel/control_v11p_sd15_shuffle",
-    "inpaint": "lllyasviel/control_v11p_sd15_inpaint",
-    "tile": "lllyasviel/control_v11p_sd15_tile",
-    "ip2p": "lllyasviel/control_v11p_sd15_ip2p",
-    "color": "lllyasviel/control_v11p_sd15_color"
+    "scribble": "lllyasviel/control_v11p_sd15_scribble",
+    "tile": "lllyasviel/control_v11f1e_sd15_tile"
 }
 
-# SDXL ControlNet models (limited availability)
+# SDXL ControlNet models
 SDXL_CONTROLNET_MODELS = {
     "canny_sdxl": "diffusers/controlnet-canny-sdxl-1.0",
-    "depth_sdxl": "diffusers/controlnet-depth-sdxl-1.0"
+    "depth_sdxl": "diffusers/controlnet-depth-sdxl-1.0",
+    "openpose_sdxl": "thibaud/controlnet-openpose-sdxl-1.0"
 }
 
-# Popular LoRA models list
+# Expanded LoRA models list (including NSFW-capable)
 LORA_MODELS = {
     "None": None,
+    # Style LoRAs
     "Lowpoly Game Character": "nerijs/lowpoly-game-character-lora",
-    "Japanese Doll": "Norod78/sd15-JapaneseDollLikeness_lora",
-    "Korean Doll": "Norod78/sd15-KoreanDollLikeness_lora",
-    "Detail Tweaker": "nitrosocke/detail-tweaker-lora",
     "Pixel Art": "nerijs/pixel-art-xl",
     "Watercolor Style": "OedoSoldier/watercolor-style-lora",
     "Manga Style": "raemikk/Animerge_V3.0_LoRA",
+    "Cyberpunk": "artificialguybr/cyberpunk-anime-diffusion",
+    "Fantasy Art": "artificialguybr/fantasy-art-lora",
+    "Chinese Style": "yfszzx/Chinese_style_xl_LoRA",
+    "Traditional Painting": "artificialguybr/Traditional-Painting-Style-LoRA",
+    "Anime Art": "Linaqruf/anime-detailer-xl-lora",
+    "Cinematic": "artificialguybr/cinematic-diffusion",
+    "Oil Painting": "artificialguybr/oil-painting-style",
+    # Character/Face LoRAs
+    "Japanese Doll": "Norod78/sd15-JapaneseDollLikeness_lora",
+    "Korean Doll": "Norod78/sd15-KoreanDollLikeness_lora",
+    "Detail Tweaker": "nitrosocke/detail-tweaker-lora",
+    "Beautiful Realistic Asians": "etok/Beautiful_Realistic_Asians",
+    "Asian Beauty": "digiplay/AsianBeauty_V1",
+    "Perfect Hands": "Sanster/perfect-hands",
+    "Face Detail": "ostris/face-detail-lora",
+    # Body/Pose LoRAs
+    "Body Pose Control": "alvdansen/lora-body-pose",
+    "Dynamic Poses": "alvdansen/dynamic-poses-lora",
+    "Full Body": "artificialguybr/full-body-lora",
+    # Realism LoRAs
     "Photorealistic": "microsoft/lora-photorealistic",
-    "Cyberpunk": "microsoft/lora-cyberpunk",
-    "Fantasy Art": "microsoft/lora-fantasy-art",
-    "Chinese Style": "yfszzx/Chinese_style_xl_LoRA",  # 中国风
-    "Traditional Painting": "artificialguybr/Traditional-Painting-Style-LoRA"
+    "Hyper-Realistic": "dallinmackay/hyper-realistic-lora",
+    "Ultra Realistic": "artificialguybr/ultra-realistic-lora",
+    "Realistic Vision": "SG161222/Realistic_Vision_V5.1_noVAE",
+    # Lighting/Quality LoRAs
+    "Add Detail": "ostris/add-detail-lora",
+    "Sharp Details": "ostris/sharp-details-lora",
+    "Better Lighting": "artificialguybr/better-lighting-lora",
+    "Studio Lighting": "artificialguybr/studio-lighting",
+    # NSFW-capable LoRAs
+    "NSFW Master": "hearmeneigh/nsfw-master-lora",
+    "Realistic NSFW": "digiplay/RealisticNSFW_v1",
+    "Anime NSFW": "Linaqruf/anime-nsfw-lora",
+    "Hentai Diffusion": "Deltaadams/Hentai-Diffusion",
+    "Sexy Pose": "alvdansen/sexy-pose-lora"
+}
+
+# VAE models for better quality
+VAE_MODELS = {
+    "None": None,
+    "SD1.5 VAE": "stabilityai/sd-vae-ft-mse",
+    "Anime VAE": "hakurei/waifu-diffusion-v1-4",
+    "SDXL VAE": "madebyollin/sdxl-vae-fp16-fix"
 }
 
 # Detector instances
 DETECTORS = {}
 
-# Florence-2 model cache
-FLORENCE2_PROCESSOR = None
-FLORENCE2_MODEL = None
-
 def is_sdxl_model(model_name: str) -> bool:
     """Check if model is SDXL"""
     return model_name in SDXL_MODELS or "xl" in model_name.lower() or "XL" in model_name
@@ -222,134 +258,6 @@ def get_controlnet_model(controlnet_type: str):
     else:
         raise ValueError(f"Unknown ControlNet type: {controlnet_type}")
 
-def load_florence2():
-    """Lazy load Florence-2 model"""
-    global FLORENCE2_PROCESSOR, FLORENCE2_MODEL
-    
-    if FLORENCE2_PROCESSOR is not None and FLORENCE2_MODEL is not None:
-        return FLORENCE2_PROCESSOR, FLORENCE2_MODEL
-    
-    try:
-        from transformers import AutoProcessor, AutoModelForCausalLM
-        
-        print("📥 Loading Microsoft/Florence-2-base...")
-        
-        # 按照官方文檔加載模型
-        FLORENCE2_MODEL = AutoModelForCausalLM.from_pretrained(
-            "microsoft/Florence-2-base", 
-            torch_dtype=torch_dtype, 
-            trust_remote_code=True
-        ).to(device)
-        
-        FLORENCE2_PROCESSOR = AutoProcessor.from_pretrained(
-            "microsoft/Florence-2-base", 
-            trust_remote_code=True
-        )
-        
-        print("✅ Florence-2 model loaded successfully")
-        return FLORENCE2_PROCESSOR, FLORENCE2_MODEL
-        
-    except Exception as e:
-        print(f"❌ Error loading Florence-2: {e}")
-        import traceback
-        traceback.print_exc()
-        return None, None
-
-def analyze_with_florence2(image, task_prompt):
-    """Analyze image using Florence-2"""
-    try:
-        processor, model = load_florence2()
-        
-        if processor is None or model is None:
-            return "❌ Failed to load Florence-2 model. Please check installation."
-        
-        # 檢查圖像
-        if image is None:
-            return "❌ No image provided for analysis."
-        
-        # 確保圖像是 PIL Image 格式
-        if not isinstance(image, Image.Image):
-            try:
-                if isinstance(image, np.ndarray):
-                    image = Image.fromarray(image)
-                else:
-                    return "❌ Invalid image format. Please upload a valid image."
-            except Exception as e:
-                return f"❌ Error converting image: {str(e)}"
-        
-        # 確保圖像是 RGB 模式
-        if image.mode != 'RGB':
-            image = image.convert('RGB')
-        
-        # 調整圖像大小以優化處理（可選）
-        max_size = 512
-        if max(image.size) > max_size:
-            ratio = max_size / max(image.size)
-            new_size = (int(image.width * ratio), int(image.height * ratio))
-            image = image.resize(new_size, Image.Resampling.LANCZOS)
-        
-        # 按照官方文檔準備輸入
-        try:
-            inputs = processor(
-                text=task_prompt,
-                images=image,
-                return_tensors="pt"
-            ).to(device, torch_dtype)
-        except Exception as e:
-            print(f"❌ Error processing image: {e}")
-            return f"❌ Error processing image: {str(e)}"
-        
-        # 按照官方文檔生成
-        try:
-            generated_ids = model.generate(
-                input_ids=inputs["input_ids"],
-                pixel_values=inputs["pixel_values"],
-                max_new_tokens=1024,
-                do_sample=False,
-                num_beams=3,
-            )
-        except Exception as e:
-            print(f"❌ Error generating text: {e}")
-            return f"❌ Error during analysis: {str(e)}"
-        
-        # 解碼
-        try:
-            generated_text = processor.batch_decode(
-                generated_ids, 
-                skip_special_tokens=False
-            )[0]
-        except Exception as e:
-            print(f"❌ Error decoding text: {e}")
-            return f"❌ Error decoding result: {str(e)}"
-        
-        # 使用 post_process_generation 解析結果
-        try:
-            parsed_answer = processor.post_process_generation(
-                generated_text, 
-                task=task_prompt,
-                image_size=(image.width, image.height)
-            )
-            
-            # 將結果轉換為可讀字符串
-            if isinstance(parsed_answer, dict):
-                result_str = ""
-                for key, value in parsed_answer.items():
-                    result_str += f"{key}:\n{value}\n\n"
-                return result_str.strip()
-            else:
-                return str(parsed_answer)
-                
-        except Exception as e:
-            print(f"❌ Error in post-processing: {e}")
-            # 如果後處理失敗，返回原始生成的文本
-            return f"Raw output: {generated_text}"
-        
-    except Exception as e:
-        print(f"❌ Error in Florence-2 analysis: {e}")
-        import traceback
-        traceback.print_exc()
-        return f"❌ Analysis error: {str(e)}"
-
 def prepare_condition_image(image, controlnet_type):
     """Prepare condition image for ControlNet"""
     if controlnet_type in ["lineart", "lineart_anime"]:
@@ -382,21 +290,19 @@ def prepare_condition_image(image, controlnet_type):
             result = detector(image, detect_resolution=512, image_resolution=512)
             return Image.fromarray(result) if isinstance(result, np.ndarray) else result
     
-    # For other types, return original image or processed version
     return image
 
-def get_pipeline(model_name: str, controlnet_type: str = "lineart", lora_model: str = None, lora_weight: float = 0.8):
-    """Get or create a ControlNet pipeline with optional LoRA"""
+def get_pipeline(model_name: str, controlnet_type: str = "lineart", lora_model: str = None, 
+                lora_weight: float = 0.8, vae_model: str = None):
+    """Get or create a ControlNet pipeline with optional LoRA and VAE"""
     global CURRENT_CONTROLNET_PIPE, CURRENT_CONTROLNET_KEY
     
-    key = (model_name, controlnet_type, lora_model, lora_weight)
+    key = (model_name, controlnet_type, lora_model, lora_weight, vae_model)
     
-    # Reuse existing pipeline
     if CURRENT_CONTROLNET_KEY == key and CURRENT_CONTROLNET_PIPE is not None:
         print(f"✅ Reusing existing ControlNet pipeline: {model_name}, type: {controlnet_type}")
         return CURRENT_CONTROLNET_PIPE
     
-    # Unload old pipeline
     if CURRENT_CONTROLNET_PIPE is not None:
         print(f"🗑️ Unloading old ControlNet pipeline: {CURRENT_CONTROLNET_KEY}")
         del CURRENT_CONTROLNET_PIPE
@@ -409,65 +315,64 @@ def get_pipeline(model_name: str, controlnet_type: str = "lineart", lora_model:
     print(f"📥 Loading ControlNet pipeline for model: {model_name}, type: {controlnet_type}")
     
     try:
-        # Check if SDXL with ControlNet
         if is_sdxl_model(model_name):
-            if controlnet_type in ["canny_sdxl", "depth_sdxl"]:
+            if controlnet_type in SDXL_CONTROLNET_MODELS:
                 controlnet_model_name = get_controlnet_model(controlnet_type)
                 controlnet = ControlNetModel.from_pretrained(
                     controlnet_model_name, 
-                    torch_dtype=torch_dtype
+                    torch_dtype=dtype
                 ).to(device)
                 
                 pipe = StableDiffusionXLPipeline.from_pretrained(
                     model_name, 
                     controlnet=controlnet, 
-                    torch_dtype=torch_dtype,
-                    safety_checker=None, 
+                    torch_dtype=dtype,
+                    safety_checker=None,
                     requires_safety_checker=False,
                     use_safetensors=True,
-                    variant="fp16" if torch_dtype == torch.float16 else None
+                    variant="fp16" if dtype == torch.float16 else None
                 ).to(device)
             else:
-                raise ValueError(f"SDXL model {model_name} only supports limited ControlNet types: {list(SDXL_CONTROLNET_MODELS.keys())}")
+                raise ValueError(f"SDXL model only supports: {list(SDXL_CONTROLNET_MODELS.keys())}")
         else:
-            # SD1.5 ControlNet
             controlnet_model_name = get_controlnet_model(controlnet_type)
             controlnet = ControlNetModel.from_pretrained(
                 controlnet_model_name, 
-                torch_dtype=torch_dtype
+                torch_dtype=dtype
             ).to(device)
             
             pipe = StableDiffusionControlNetPipeline.from_pretrained(
                 model_name, 
                 controlnet=controlnet, 
-                torch_dtype=torch_dtype,
-                safety_checker=None, 
+                torch_dtype=dtype,
+                safety_checker=None,
                 requires_safety_checker=False,
                 use_safetensors=True,
-                variant="fp16" if torch_dtype == torch.float16 else None
+                variant="fp16" if dtype == torch.float16 else None
             ).to(device)
         
+        # Load custom VAE if specified
+        if vae_model and vae_model != "None":
+            try:
+                from diffusers import AutoencoderKL
+                print(f"🔄 Loading custom VAE: {vae_model}")
+                vae = AutoencoderKL.from_pretrained(vae_model, torch_dtype=dtype).to(device)
+                pipe.vae = vae
+            except Exception as e:
+                print(f"⚠️  Error loading VAE: {e}")
+        
         # Apply LoRA if specified
         if lora_model and lora_model != "None":
             print(f"🔄 Applying LoRA: {lora_model} with weight: {lora_weight}")
             try:
-                pipe.load_lora_weights(lora_model, weight_name=None if "safetensors" in lora_model else "pytorch_lora_weights.safetensors")
+                pipe.load_lora_weights(lora_model)
                 pipe.fuse_lora(lora_scale=lora_weight)
             except Exception as e:
                 print(f"⚠️  Error loading LoRA: {e}")
-                print("Trying alternative LoRA loading method...")
-                try:
-                    from safetensors.torch import load_file
-                    from huggingface_hub import hf_hub_download
-                    lora_path = hf_hub_download(lora_model, "pytorch_lora_weights.safetensors")
-                    pipe.unet.load_state_dict(load_file(lora_path), strict=False)
-                except Exception as e2:
-                    print(f"❌ Failed to load LoRA: {e2}")
         
         # Optimizations
         pipe.enable_attention_slicing(slice_size="max")
         
-        # VAE slicing
         if hasattr(pipe, 'vae') and hasattr(pipe.vae, 'enable_slicing'):
             pipe.vae.enable_slicing()
         else:
@@ -477,34 +382,19 @@ def get_pipeline(model_name: str, controlnet_type: str = "lineart", lora_model:
                 pass
         
         if device.type == "cuda":
-            # xFormers
             try:
                 pipe.enable_xformers_memory_efficient_attention()
-                print("✅ xFormers enabled for ControlNet")
+                print("✅ xFormers enabled")
             except:
-                print("⚠️  xFormers not available, using standard attention")
                 pass
-            
-            # Model CPU offload
             pipe.enable_model_cpu_offload()
         
-        # Compile model for faster inference
-        if hasattr(torch, 'compile') and device.type == "cuda":
-            try:
-                pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-                print("✅ Model compiled with torch.compile")
-            except Exception as e:
-                print(f"⚠️  torch.compile not available: {e}")
-                pass
-        
-        # Change scheduler for better quality
         try:
-            pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
-            print("✅ Using UniPC scheduler for faster convergence")
+            pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+            print("✅ Using Euler Ancestral scheduler")
         except:
             try:
                 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-                print("✅ Using DPM++ scheduler")
             except:
                 pass
         
@@ -518,19 +408,17 @@ def get_pipeline(model_name: str, controlnet_type: str = "lineart", lora_model:
         CURRENT_CONTROLNET_KEY = None
         raise
 
-def load_t2i_model(model_name: str, lora_model: str = None, lora_weight: float = 0.8):
-    """Load text-to-image model with optional LoRA"""
+def load_t2i_model(model_name: str, lora_model: str = None, lora_weight: float = 0.8, 
+                   vae_model: str = None):
+    """Load text-to-image model with optional LoRA and VAE"""
     global CURRENT_T2I_PIPE, CURRENT_T2I_MODEL, CURRENT_SDXL_REFINER
     
-    # Check if we need to load refiner for SDXL
     use_refiner = "refiner" in model_name.lower()
-    
-    key = (model_name, lora_model, lora_weight, use_refiner)
+    key = (model_name, lora_model, lora_weight, vae_model, use_refiner)
     
     if CURRENT_T2I_MODEL == key and CURRENT_T2I_PIPE is not None:
         return
     
-    # Unload old model
     if CURRENT_T2I_PIPE is not None:
         print(f"🗑️ Unloading old T2I model: {CURRENT_T2I_MODEL}")
         del CURRENT_T2I_PIPE
@@ -546,64 +434,67 @@ def load_t2i_model(model_name: str, lora_model: str = None, lora_weight: float =
     
     try:
         if is_sdxl_model(model_name):
-            # Load SDXL model
             if use_refiner:
-                # Load base and refiner
                 CURRENT_T2I_PIPE = StableDiffusionXLPipeline.from_pretrained(
                     "stabilityai/stable-diffusion-xl-base-1.0",
-                    torch_dtype=torch_dtype, 
-                    safety_checker=None, 
+                    torch_dtype=dtype,
+                    safety_checker=None,
                     requires_safety_checker=False,
                     use_safetensors=True,
-                    variant="fp16" if torch_dtype == torch.float16 else None
+                    variant="fp16" if dtype == torch.float16 else None
                 ).to(device)
                 
                 CURRENT_SDXL_REFINER = StableDiffusionXLPipeline.from_pretrained(
                     model_name,
-                    torch_dtype=torch_dtype,
+                    torch_dtype=dtype,
                     safety_checker=None,
                     requires_safety_checker=False,
                     use_safetensors=True,
-                    variant="fp16" if torch_dtype == torch.float16 else None,
+                    variant="fp16" if dtype == torch.float16 else None,
                     text_encoder_2=CURRENT_T2I_PIPE.text_encoder_2,
                     vae=CURRENT_T2I_PIPE.vae
                 ).to(device)
-                print(f"✅ Loaded SDXL with refiner: {model_name}")
             else:
                 CURRENT_T2I_PIPE = StableDiffusionXLPipeline.from_pretrained(
-                    model_name, 
-                    torch_dtype=torch_dtype, 
-                    safety_checker=None, 
+                    model_name,
+                    torch_dtype=dtype,
+                    safety_checker=None,
                     requires_safety_checker=False,
                     use_safetensors=True,
-                    variant="fp16" if torch_dtype == torch.float16 else None
+                    variant="fp16" if dtype == torch.float16 else None
                 ).to(device)
-                print(f"✅ Loaded SDXL model: {model_name}")
         else:
-            # Load SD1.5 model
             CURRENT_T2I_PIPE = StableDiffusionPipeline.from_pretrained(
-                model_name, 
-                torch_dtype=torch_dtype, 
-                safety_checker=None, 
+                model_name,
+                torch_dtype=dtype,
+                safety_checker=None,
                 requires_safety_checker=False,
                 use_safetensors=True,
-                variant="fp16" if torch_dtype == torch.float16 else None
+                variant="fp16" if dtype == torch.float16 else None
             ).to(device)
-            print(f"✅ Loaded SD1.5 model: {model_name}")
         
-        # Apply LoRA if specified
+        # Load custom VAE
+        if vae_model and vae_model != "None":
+            try:
+                from diffusers import AutoencoderKL
+                print(f"🔄 Loading custom VAE: {vae_model}")
+                vae = AutoencoderKL.from_pretrained(vae_model, torch_dtype=dtype).to(device)
+                CURRENT_T2I_PIPE.vae = vae
+            except Exception as e:
+                print(f"⚠️  Error loading VAE: {e}")
+        
+        # Apply LoRA
         if lora_model and lora_model != "None":
-            print(f"🔄 Applying LoRA to T2I: {lora_model} with weight: {lora_weight}")
+            print(f"🔄 Applying LoRA: {lora_model} with weight: {lora_weight}")
             try:
                 CURRENT_T2I_PIPE.load_lora_weights(lora_model)
                 CURRENT_T2I_PIPE.fuse_lora(lora_scale=lora_weight)
             except Exception as e:
-                print(f"⚠️  Error loading LoRA for T2I: {e}")
+                print(f"⚠️  Error loading LoRA: {e}")
         
         # Optimizations
         CURRENT_T2I_PIPE.enable_attention_slicing(slice_size="max")
         
-        # VAE slicing
         if hasattr(CURRENT_T2I_PIPE, 'vae') and hasattr(CURRENT_T2I_PIPE.vae, 'enable_slicing'):
             CURRENT_T2I_PIPE.vae.enable_slicing()
         else:
@@ -615,129 +506,26 @@ def load_t2i_model(model_name: str, lora_model: str = None, lora_weight: float =
         if device.type == "cuda":
             try:
                 CURRENT_T2I_PIPE.enable_xformers_memory_efficient_attention()
-                print("✅ xFormers enabled for T2I")
             except:
                 pass
             CURRENT_T2I_PIPE.enable_model_cpu_offload()
         
-        # Change scheduler
         try:
-            CURRENT_T2I_PIPE.scheduler = UniPCMultistepScheduler.from_config(CURRENT_T2I_PIPE.scheduler.config)
-            print("✅ Using UniPC scheduler")
+            CURRENT_T2I_PIPE.scheduler = EulerAncestralDiscreteScheduler.from_config(CURRENT_T2I_PIPE.scheduler.config)
         except:
-            try:
-                CURRENT_T2I_PIPE.scheduler = DPMSolverMultistepScheduler.from_config(CURRENT_T2I_PIPE.scheduler.config)
-                print("✅ Using DPM++ scheduler")
-            except:
-                pass
+            pass
         
         CURRENT_T2I_MODEL = key
         
     except Exception as e:
-        print(f"❌ Error loading T2I model {model_name}: {e}")
-        print(f"⚠️  Trying to load without use_safetensors...")
-        
-        # Retry without use_safetensors
-        try:
-            if is_sdxl_model(model_name):
-                CURRENT_T2I_PIPE = StableDiffusionXLPipeline.from_pretrained(
-                    model_name, 
-                    torch_dtype=torch_dtype, 
-                    safety_checker=None, 
-                    requires_safety_checker=False
-                ).to(device)
-            else:
-                CURRENT_T2I_PIPE = StableDiffusionPipeline.from_pretrained(
-                    model_name, 
-                    torch_dtype=torch_dtype, 
-                    safety_checker=None, 
-                    requires_safety_checker=False
-                ).to(device)
-            
-            # Optimizations
-            CURRENT_T2I_PIPE.enable_attention_slicing(slice_size="max")
-            if hasattr(CURRENT_T2I_PIPE, 'vae') and hasattr(CURRENT_T2I_PIPE.vae, 'enable_slicing'):
-                CURRENT_T2I_PIPE.vae.enable_slicing()
-            else:
-                try:
-                    CURRENT_T2I_PIPE.enable_vae_slicing()
-                except:
-                    pass
-            
-            if device.type == "cuda":
-                try:
-                    CURRENT_T2I_PIPE.enable_xformers_memory_efficient_attention()
-                    print("✅ xFormers enabled for T2I")
-                except:
-                    pass
-                CURRENT_T2I_PIPE.enable_model_cpu_offload()
-            
-            CURRENT_T2I_MODEL = key
-            
-        except Exception as retry_e:
-            print(f"❌ Error loading T2I model (retry): {retry_e}")
-            CURRENT_T2I_PIPE = None
-            CURRENT_T2I_MODEL = None
-            raise
-
-# ===== Utils =====
-def resize_image(image, max_size=1024):
-    """Resize image while maintaining aspect ratio"""
-    width, height = image.size
-    if max(width, height) > max_size:
-        ratio = max_size / max(width, height)
-        new_width = int(width * ratio)
-        new_height = int(height * ratio)
-        return image.resize((new_width, new_height), Image.LANCZOS)
-    return image
-
-def image_to_image(img, prompt, negative_prompt, model_name, strength=0.75, steps=30, scale=7.5, seed=42):
-    """Image-to-Image transformation"""
-    try:
-        load_t2i_model(model_name)
-        
-        # Resize if needed
-        img = resize_image(img, 1024)
-        
-        # Create img2img pipeline
-        pipe = StableDiffusionImg2ImgPipeline(
-            vae=CURRENT_T2I_PIPE.vae,
-            text_encoder=CURRENT_T2I_PIPE.text_encoder,
-            tokenizer=CURRENT_T2I_PIPE.tokenizer,
-            unet=CURRENT_T2I_PIPE.unet,
-            scheduler=CURRENT_T2I_PIPE.scheduler,
-            safety_checker=None,
-            feature_extractor=None,
-            requires_safety_checker=False,
-        ).to(device)
-        
-        gen = torch.Generator(device=device).manual_seed(int(seed))
-        
-        with torch.inference_mode():
-            result = pipe(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                image=img,
-                strength=strength,
-                num_inference_steps=int(steps),
-                guidance_scale=float(scale),
-                generator=gen
-            ).images[0]
-        
-        if device.type == "cuda":
-            torch.cuda.empty_cache()
-        
-        return result
-    except Exception as e:
-        print(f"❌ Error in img2img: {e}")
-        error_img = Image.new('RGB', (512, 512), color='red')
-        return error_img
+        print(f"❌ Error loading T2I model: {e}")
+        CURRENT_T2I_PIPE = None
+        CURRENT_T2I_MODEL = None
+        raise
 
-# ===== Functions =====
-def colorize(sketch, base_model, controlnet_type, lora_model, lora_weight, 
+def colorize(sketch, base_model, controlnet_type, lora_model, lora_weight, vae_model,
              prompt, negative_prompt, seed, steps, scale, cn_weight):
     try:
-        # 檢查是否為 SDXL model 且不支援 ControlNet
         if is_sdxl_model(base_model) and controlnet_type not in SDXL_CONTROLNET_MODELS:
             error_img = Image.new('RGB', (512, 512), color='red')
             error_msg_img = Image.new('RGB', (512, 512), color='yellow')
@@ -751,34 +539,30 @@ def colorize(sketch, base_model, controlnet_type, lora_model, lora_weight,
             draw.text((50, 230), f"{', '.join(SDXL_CONTROLNET_MODELS.keys())}", fill="black", font=font)
             return error_img, error_msg_img
         
-        # 載入 pipeline
-        pipe = get_pipeline(base_model, controlnet_type, lora_model, lora_weight)
+        pipe = get_pipeline(base_model, controlnet_type, lora_model, lora_weight, vae_model)
         
-        status_msg = f"🎨 Using: {base_model} + {controlnet_type} ControlNet"
+        status_msg = f"🎨 Using: {base_model} + {controlnet_type}"
         if lora_model and lora_model != "None":
             status_msg += f" + {lora_model}"
         print(status_msg)
         
-        # 準備 condition image
         condition_img = prepare_condition_image(sketch, controlnet_type)
         
-        # 生成圖像
         gen = torch.Generator(device=device).manual_seed(int(seed))
         
         with torch.inference_mode():
             out = pipe(
-                prompt, 
+                prompt,
                 negative_prompt=negative_prompt,
-                image=condition_img, 
+                image=condition_img,
                 num_inference_steps=int(steps),
-                guidance_scale=float(scale), 
+                guidance_scale=float(scale),
                 controlnet_conditioning_scale=float(cn_weight),
                 generator=gen,
                 height=512,
                 width=512
             ).images[0]
         
-        # Clear cache
         if device.type == "cuda":
             torch.cuda.empty_cache()
         
@@ -788,14 +572,14 @@ def colorize(sketch, base_model, controlnet_type, lora_model, lora_weight,
         error_img = Image.new('RGB', (512, 512), color='red')
         return error_img, Image.new('RGB', (512, 512), color='gray')
 
-def t2i(prompt, negative_prompt, model, lora_model, lora_weight, seed, steps, scale, w, h, use_refiner=False):
+def t2i(prompt, negative_prompt, model, lora_model, lora_weight, vae_model,
+        seed, steps, scale, w, h, use_refiner=False):
     try:
-        # 如果需要 refiner，使用特殊的模型名稱
         model_to_load = model
         if use_refiner and "refiner" not in model.lower():
             model_to_load = "stabilityai/stable-diffusion-xl-refiner-1.0"
         
-        load_t2i_model(model_to_load, lora_model, lora_weight)
+        load_t2i_model(model_to_load, lora_model, lora_weight, vae_model)
         
         print(f"🖼️ Using T2I model: {model}")
         if lora_model and lora_model != "None":
@@ -804,50 +588,46 @@ def t2i(prompt, negative_prompt, model, lora_model, lora_weight, seed, steps, sc
         gen = torch.Generator(device=device).manual_seed(int(seed))
         
         with torch.inference_mode():
-            # SDXL with refiner
             if use_refiner and CURRENT_SDXL_REFINER is not None:
-                # First stage with base model
                 image = CURRENT_T2I_PIPE(
                     prompt=prompt,
                     negative_prompt=negative_prompt,
                     width=int(w),
                     height=int(h),
-                    num_inference_steps=int(steps//2),  # Half steps for base
+                    num_inference_steps=int(steps//2),
                     guidance_scale=float(scale),
                     generator=gen,
                     output_type="latent"
                 ).images
                 
-                # Second stage with refiner
                 result = CURRENT_SDXL_REFINER(
                     prompt=prompt,
                     negative_prompt=negative_prompt,
                     image=image,
-                    num_inference_steps=int(steps//2),  # Half steps for refiner
+                    num_inference_steps=int(steps//2),
                     guidance_scale=float(scale),
                     generator=gen
                 ).images[0]
             else:
-                # Normal generation
                 if is_sdxl_model(model):
                     width = max(int(w), 512)
                     height = max(int(h), 512)
                     result = CURRENT_T2I_PIPE(
-                        prompt, 
+                        prompt,
                         negative_prompt=negative_prompt,
-                        width=width, 
+                        width=width,
                         height=height,
-                        num_inference_steps=int(steps), 
+                        num_inference_steps=int(steps),
                         guidance_scale=float(scale),
                         generator=gen
                     ).images[0]
                 else:
                     result = CURRENT_T2I_PIPE(
-                        prompt, 
+                        prompt,
                         negative_prompt=negative_prompt,
-                        width=int(w), 
+                        width=int(w),
                         height=int(h),
-                        num_inference_steps=int(steps), 
+                        num_inference_steps=int(steps),
                         guidance_scale=float(scale),
                         generator=gen
                     ).images[0]
@@ -868,39 +648,13 @@ def t2i(prompt, negative_prompt, model, lora_model, lora_weight, seed, steps, sc
         draw.text((50, 50), f"Error: {str(e)[:50]}...", fill="white", font=font)
         return error_img
 
-def florence2_analysis(image, task_prompt, custom_prompt):
-    """Analyze image with Florence-2"""
-    try:
-        if image is None:
-            return "❌ Please upload an image first"
-        
-        # 確保圖像是 PIL Image 格式
-        if not isinstance(image, Image.Image):
-            return "❌ Invalid image format. Please upload a valid image."
-        
-        # Use custom prompt if provided
-        prompt_to_use = custom_prompt.strip() if custom_prompt.strip() else task_prompt
-        
-        print(f"🔍 Analyzing image with Florence-2 using prompt: {prompt_to_use}")
-        result = analyze_with_florence2(image, prompt_to_use)
-        return result
-        
-    except Exception as e:
-        print(f"❌ Error in Florence-2 analysis: {e}")
-        import traceback
-        traceback.print_exc()
-        return f"Error: {str(e)}"
-
-# ===== Function to unload all models =====
 def unload_all_models():
     global CURRENT_CONTROLNET_PIPE, CURRENT_CONTROLNET_KEY
     global DETECTORS
     global CURRENT_T2I_PIPE, CURRENT_T2I_MODEL, CURRENT_SDXL_REFINER
-    global FLORENCE2_PROCESSOR, FLORENCE2_MODEL
     
-    print("Unloading all models from memory...")
+    print("🗑️ Unloading all models from memory...")
     
-    # Unload ControlNet pipeline
     try:
         if CURRENT_CONTROLNET_PIPE is not None:
             del CURRENT_CONTROLNET_PIPE
@@ -909,7 +663,6 @@ def unload_all_models():
         pass
     CURRENT_CONTROLNET_KEY = None
     
-    # Unload detectors
     for detector_type in list(DETECTORS.keys()):
         try:
             del DETECTORS[detector_type]
@@ -917,7 +670,6 @@ def unload_all_models():
             pass
     DETECTORS.clear()
     
-    # Unload T2I models
     try:
         if CURRENT_T2I_PIPE is not None:
             del CURRENT_T2I_PIPE
@@ -934,22 +686,6 @@ def unload_all_models():
     
     CURRENT_T2I_MODEL = None
     
-    # Unload Florence-2
-    try:
-        if FLORENCE2_PROCESSOR is not None:
-            del FLORENCE2_PROCESSOR
-            FLORENCE2_PROCESSOR = None
-    except:
-        pass
-    
-    try:
-        if FLORENCE2_MODEL is not None:
-            del FLORENCE2_MODEL
-            FLORENCE2_MODEL = None
-    except:
-        pass
-    
-    # Force garbage collection
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
@@ -960,12 +696,11 @@ def unload_all_models():
     return "✅ All models unloaded from memory!"
 
 # ===== Gradio UI =====
-with gr.Blocks(title="🎨 Advanced Image Generation Suite", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🎨 Advanced Image Generation & Editing Suite")
-    gr.Markdown("### Powered by Stable Diffusion & ControlNet")
-    gr.Markdown("**Note:** SDXL models work with limited ControlNet types (canny_sdxl, depth_sdxl)")
+with gr.Blocks(title="🎨 AI Image Generator Pro", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🎨 AI Image Generator Pro - NSFW Capable")
+    gr.Markdown("### Advanced Image Generation with ControlNet, LoRA & VAE Support")
+    gr.Markdown("⚠️ **Content Warning:** This tool can generate NSFW content. Use responsibly and in compliance with applicable laws.")
     
-    # System info
     if torch.cuda.is_available():
         gpu_name = torch.cuda.get_device_name(0)
         gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
@@ -973,301 +708,257 @@ with gr.Blocks(title="🎨 Advanced Image Generation Suite", theme=gr.themes.Sof
     else:
         gr.Markdown("**⚠️ Running on CPU** - Generation will be slower")
     
-    # Add unload button
     with gr.Row():
         unload_btn = gr.Button("🗑️ Unload All Models", variant="stop", scale=1)
         status_text = gr.Textbox(label="Status", interactive=False, scale=3)
     unload_btn.click(unload_all_models, outputs=status_text)
 
-    with gr.Tab("🎨 ControlNet Colorize"):
+    with gr.Tab("🎨 ControlNet Image-to-Image"):
         gr.Markdown("""
-        ### Convert sketches to colored images using ControlNet
-        **SD1.5 Models:** Support all ControlNet types
-        **SDXL Models:** Only support canny_sdxl and depth_sdxl
+        ### Transform sketches/images using ControlNet
+        - **SD1.5 Models:** Support all ControlNet types
+        - **SDXL Models:** Support canny_sdxl, depth_sdxl, openpose_sdxl only
         """)
         
         with gr.Row():
-            inp = gr.Image(label="Input Sketch/Image", type="pil")
-            out = gr.Image(label="Colored Output")
-        
-        with gr.Row():
-            condition_out = gr.Image(label="Processed Condition Image", type="pil")
-        
-        with gr.Row():
-            base_model = gr.Dropdown(
-                choices=ALL_MODELS,
-                value="digiplay/ChikMix_V3",
-                label="Base Model"
-            )
-            controlnet_type = gr.Dropdown(
-                choices=list(CONTROLNET_MODELS.keys()) + list(SDXL_CONTROLNET_MODELS.keys()),
-                value="lineart_anime",
-                label="ControlNet Type"
-            )
-        
-        with gr.Row():
-            lora_model = gr.Dropdown(
-                choices=list(LORA_MODELS.keys()),
-                value="None",
-                label="LoRA Model (Optional)"
-            )
-            lora_weight = gr.Slider(0.1, 1.5, 0.8, step=0.1, label="LoRA Weight")
+            with gr.Column(scale=1):
+                inp = gr.Image(label="Input Sketch/Image", type="pil")
+                
+                gr.Markdown("### Model Settings")
+                base_model = gr.Dropdown(
+                    choices=ALL_MODELS,
+                    value="digiplay/ChikMix_V3",
+                    label="Base Model"
+                )
+                controlnet_type = gr.Dropdown(
+                    choices=list(CONTROLNET_MODELS.keys()) + list(SDXL_CONTROLNET_MODELS.keys()),
+                    value="lineart_anime",
+                    label="ControlNet Type"
+                )
+                
+                gr.Markdown("### Enhancement Options")
+                with gr.Row():
+                    lora_model = gr.Dropdown(
+                        choices=list(LORA_MODELS.keys()),
+                        value="None",
+                        label="LoRA Model"
+                    )
+                    lora_weight = gr.Slider(0.1, 2.0, 0.8, step=0.1, label="LoRA Weight")
+                
+                vae_model = gr.Dropdown(
+                    choices=list(VAE_MODELS.keys()),
+                    value="None",
+                    label="VAE Model (Optional)"
+                )
+            
+            with gr.Column(scale=1):
+                out = gr.Image(label="Generated Output")
+                condition_out = gr.Image(label="Processed Condition", type="pil")
         
+        gr.Markdown("### Generation Parameters")
         with gr.Row():
             prompt = gr.Textbox(
-                label="Prompt", 
-                placeholder="e.g., 1girl, blonde hair, blue eyes, beautiful, masterpiece",
-                lines=2
+                label="Prompt",
+                placeholder="masterpiece, best quality, 1girl, beautiful detailed eyes, long hair",
+                lines=3
             )
             negative_prompt = gr.Textbox(
-                label="Negative Prompt", 
-                placeholder="e.g., ugly, deformed, bad anatomy, blurry",
-                lines=2
+                label="Negative Prompt",
+                placeholder="lowres, bad anatomy, bad hands, text, error, missing fingers",
+                lines=3
             )
         
         with gr.Row():
-            seed = gr.Number(value=42, label="Seed")
-            steps = gr.Slider(10, 100, 30, step=1, label="Steps")
-            scale = gr.Slider(1, 20, 7.5, step=0.5, label="CFG Scale")
+            seed = gr.Number(value=-1, label="Seed (-1 for random)")
+            steps = gr.Slider(10, 150, 30, step=1, label="Steps")
+            scale = gr.Slider(1, 30, 7.5, step=0.5, label="CFG Scale")
             cn_weight = gr.Slider(0.1, 2.0, 1.0, step=0.1, label="ControlNet Weight")
         
-        run = gr.Button("🎨 Colorize", variant="primary")
+        run = gr.Button("🎨 Generate", variant="primary", size="lg")
         run.click(
-            colorize, 
-            [inp, base_model, controlnet_type, lora_model, lora_weight, 
-             prompt, negative_prompt, seed, steps, scale, cn_weight], 
+            colorize,
+            [inp, base_model, controlnet_type, lora_model, lora_weight, vae_model,
+             prompt, negative_prompt, seed, steps, scale, cn_weight],
             [out, condition_out]
         )
+        
+        gr.Markdown("""
+        ### Tips for Better Results:
+        - Use detailed prompts for better control
+        - Adjust ControlNet weight to balance between condition and creativity
+        - Try different LoRA models for various styles
+        - Higher steps = better quality but slower generation
+        """)
 
-    with gr.Tab("🖼️ Text-to-Image"):
+    with gr.Tab("🖼️ Text-to-Image Generation"):
         gr.Markdown("""
         ### Generate images from text descriptions
-        Supports both SD1.5 and SDXL models with optional LoRA.
-        **Tip:** SDXL models produce higher quality but require more memory.
+        Supports both SD1.5 and SDXL models with advanced features
         """)
         
         with gr.Row():
-            t2i_out = gr.Image(label="Output", type="pil")
+            with gr.Column(scale=1):
+                gr.Markdown("### Model Configuration")
+                t2i_model = gr.Dropdown(
+                    choices=ALL_MODELS,
+                    value="digiplay/ChikMix_V3",
+                    label="Base Model"
+                )
+                
+                gr.Markdown("### Enhancement Options")
+                with gr.Row():
+                    t2i_lora = gr.Dropdown(
+                        choices=list(LORA_MODELS.keys()),
+                        value="None",
+                        label="LoRA Model"
+                    )
+                    t2i_lora_weight = gr.Slider(0.1, 2.0, 0.8, step=0.1, label="LoRA Weight")
+                
+                t2i_vae = gr.Dropdown(
+                    choices=list(VAE_MODELS.keys()),
+                    value="None",
+                    label="VAE Model"
+                )
+                
+                use_refiner = gr.Checkbox(
+                    label="Use SDXL Refiner (SDXL only)",
+                    value=False
+                )
+            
+            with gr.Column(scale=1):
+                t2i_out = gr.Image(label="Generated Image", type="pil")
         
+        gr.Markdown("### Prompts")
         with gr.Row():
             t2i_prompt = gr.Textbox(
-                label="Prompt", 
-                lines=3,
-                placeholder="e.g., a beautiful landscape with mountains and a lake at sunset, highly detailed, 4k"
+                label="Prompt",
+                lines=4,
+                placeholder="masterpiece, best quality, highly detailed, 8k, photorealistic, beautiful lighting"
             )
             t2i_negative_prompt = gr.Textbox(
-                label="Negative Prompt", 
-                lines=2,
-                placeholder="e.g., blurry, ugly, deformed, low quality"
-            )
-        
-        with gr.Row():
-            t2i_model = gr.Dropdown(
-                choices=ALL_MODELS,
-                value="digiplay/ChikMix_V3",
-                label="Model"
-            )
-            t2i_lora = gr.Dropdown(
-                choices=list(LORA_MODELS.keys()),
-                value="None",
-                label="LoRA Model (Optional)"
+                label="Negative Prompt",
+                lines=4,
+                placeholder="lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality"
             )
-            t2i_lora_weight = gr.Slider(0.1, 1.5, 0.8, step=0.1, label="LoRA Weight")
         
+        gr.Markdown("### Generation Parameters")
         with gr.Row():
-            t2i_seed = gr.Number(value=42, label="Seed")
-            t2i_steps = gr.Slider(10, 100, 30, step=1, label="Steps")
-            t2i_scale = gr.Slider(1, 20, 7.5, step=0.5, label="CFG Scale")
+            t2i_seed = gr.Number(value=-1, label="Seed (-1 for random)")
+            t2i_steps = gr.Slider(10, 150, 30, step=1, label="Steps")
+            t2i_scale = gr.Slider(1, 30, 7.5, step=0.5, label="CFG Scale")
         
         with gr.Row():
-            w = gr.Slider(256, 2048, 1024, step=64, label="Width")
-            h = gr.Slider(256, 2048, 1024, step=64, label="Height")
-            use_refiner = gr.Checkbox(label="Use SDXL Refiner (SDXL only)", value=False)
+            w = gr.Slider(256, 2048, 512, step=64, label="Width")
+            h = gr.Slider(256, 2048, 768, step=64, label="Height")
         
-        gen_btn = gr.Button("🖼️ Generate", variant="primary")
+        gen_btn = gr.Button("🖼️ Generate Image", variant="primary", size="lg")
         gen_btn.click(
-            t2i, 
-            [t2i_prompt, t2i_negative_prompt, t2i_model, t2i_lora, t2i_lora_weight, 
-             t2i_seed, t2i_steps, t2i_scale, w, h, use_refiner], 
+            t2i,
+            [t2i_prompt, t2i_negative_prompt, t2i_model, t2i_lora, t2i_lora_weight,
+             t2i_vae, t2i_seed, t2i_steps, t2i_scale, w, h, use_refiner],
             t2i_out
         )
-
-    with gr.Tab("🔄 Image-to-Image"):
-        gr.Markdown("""
-        ### Transform existing images using img2img
-        Modify images based on prompts with control over transformation strength.
-        """)
-        
-        with gr.Row():
-            img2img_input = gr.Image(label="Input Image", type="pil")
-            img2img_output = gr.Image(label="Transformed Output")
-        
-        with gr.Row():
-            img2img_prompt = gr.Textbox(
-                label="Prompt",
-                lines=2,
-                placeholder="e.g., make it anime style, cyberpunk style, etc."
-            )
-            img2img_negative_prompt = gr.Textbox(
-                label="Negative Prompt",
-                lines=2,
-                placeholder="e.g., blurry, low quality"
-            )
         
-        with gr.Row():
-            img2img_model = gr.Dropdown(
-                choices=ALL_MODELS,
-                value="stablediffusionapi/realistic-vision-v51",
-                label="Model"
-            )
-            img2img_strength = gr.Slider(0.1, 0.95, 0.75, step=0.05, label="Transformation Strength")
-        
-        with gr.Row():
-            img2img_seed = gr.Number(value=42, label="Seed")
-            img2img_steps = gr.Slider(10, 100, 30, step=1, label="Steps")
-            img2img_scale = gr.Slider(1, 20, 7.5, step=0.5, label="CFG Scale")
-        
-        img2img_btn = gr.Button("🔄 Transform Image", variant="primary")
-        img2img_btn.click(
-            image_to_image,
-            [img2img_input, img2img_prompt, img2img_negative_prompt, 
-             img2img_model, img2img_strength, img2img_steps, img2img_scale, img2img_seed],
-            img2img_output
-        )
-
-    with gr.Tab("🔍 Florence-2 Vision Analysis"):
         gr.Markdown("""
-        ### Microsoft Florence-2 Vision Language Model
-        **Pre-trained Tasks:**
-        - `<OCR>`: Text recognition (Extract text from image)
-        - `<CAPTION>`: Image captioning (Generate a caption)
-        - `<DETAILED_CAPTION>`: Detailed caption (More detailed description)
-        - `<MORE_DETAILED_CAPTION>`: More detailed caption (Even more details)
-        - `<OD>`: Object detection (Detect objects with bounding boxes)
-        - `<OPEN_VOCABULARY_DETECTION>`: Open-vocabulary detection
-        - `<REGION_PROPOSAL>`: Region proposal
-        
-        **How to use:**
-        1. Upload an image
-        2. Select a task from the dropdown
-        3. Click "Analyze Image"
-        4. Results will be displayed in the text box
-        
-        **Example tasks:**
-        - Extract text from a document: `<OCR>`
-        - Describe what's in the image: `<CAPTION>`
-        - Detect objects in the image: `<OD>`
+        ### Pro Tips:
+        - **SDXL models** produce higher quality at 1024x1024
+        - **SD1.5 models** work best at 512x512 or 512x768
+        - Use **LoRA** for specific styles (anime, realistic, etc.)
+        - Use **VAE** for better colors and details
+        - **Refiner** adds extra polish to SDXL generations
+        - Higher **CFG Scale** = more prompt adherence
         """)
-        
-        with gr.Row():
-            florence_input = gr.Image(label="Input Image", type="pil")
-            florence_output = gr.Textbox(
-                label="Analysis Result", 
-                lines=15,
-                interactive=False,
-                show_copy_button=True
-            )
-        
-        with gr.Row():
-            florence_task = gr.Dropdown(
-                choices=[
-                    "<OCR>",
-                    "<CAPTION>",
-                    "<DETAILED_CAPTION>",
-                    "<MORE_DETAILED_CAPTION>",
-                    "<OD>",
-                    "<OPEN_VOCABULARY_DETECTION>",
-                    "<REGION_PROPOSAL>"
-                ],
-                value="<CAPTION>",
-                label="Task Prompt"
-            )
-            
-            custom_prompt = gr.Textbox(
-                label="Custom Prompt (Optional)",
-                value="",
-                placeholder="e.g., Describe the main objects in this image"
-            )
-        
-        with gr.Row():
-            analyze_btn = gr.Button("🔍 Analyze Image", variant="primary")
-            clear_btn = gr.Button("🗑️ Clear")
-        
-        def clear_analysis():
-            return None, ""
-        
-        analyze_btn.click(
-            florence2_analysis,
-            [florence_input, florence_task, custom_prompt],
-            florence_output
-        )
-        
-        clear_btn.click(
-            clear_analysis,
-            [],
-            [florence_input, florence_output]
-        )
-
-    with gr.Tab("📊 Model Info"):
+    
+    with gr.Tab("📚 Quick Reference"):
         gr.Markdown("""
-        ### Available Models Information
-        
-        **SD1.5 Models (Support all ControlNet types):**
-        - Recommended for ControlNet workflows
-        - Faster inference, lower memory usage
-        - Wide variety of styles available
-        
-        **SDXL Models (Higher quality, limited ControlNet):**
-        - Better quality, more details
-        - Larger image sizes (1024x1024+)
-        - Only supports canny_sdxl and depth_sdxl ControlNet
-        
-        **Chinese Models:**
-        - Optimized for Chinese prompts
-        - Better understanding of Chinese culture elements
-        
-        **Florence-2 Model:**
-        - Microsoft's vision-language model
-        - Image analysis, OCR, captioning, object detection
-        
-        **LoRA Models:**
-        - Fine-tuned models for specific styles
-        - Can be combined with base models
-        - Adjust weight for stronger/weaker effect
+        # Model & Feature Guide
+        
+        ## 🎯 Recommended Models for Different Purposes
+        
+        ### Realistic/Photorealistic
+        - `emilianJR/epiCRealism` - Excellent for realistic portraits
+        - `stablediffusionapi/realistic-vision-v51` - High quality realistic images
+        - `digiplay/majicMIX_realistic_v7` - Great for realistic characters
+        - `SG161222/RealVisXL_V4.0` - SDXL realistic model
+        
+        ### Anime/Cartoon
+        - `digiplay/ChikMix_V3` - Versatile anime style
+        - `gsdf/Counterfeit-V2.5` - High quality anime
+        - `stablediffusionapi/anything-v5` - Popular anime model
+        - `digiplay/Pony_Diffusion_V6_XL` - SDXL anime model
+        
+        ### Artistic/Stylized
+        - `stablediffusionapi/dreamshaper-v8` - Dream-like artistic style
+        - `wavymulder/Analog-Diffusion` - Analog photo aesthetic
+        - `Lykon/dreamshaper-xl-1-0` - SDXL artistic model
+        
+        ## 🎨 ControlNet Types Explained
+        
+        - **lineart/lineart_anime**: Convert line drawings to colored images
+        - **canny**: Edge detection based generation
+        - **depth**: Depth map based generation
+        - **openpose**: Human pose based generation
+        - **normal**: Normal map based generation
+        - **softedge**: Soft edge detection
+        - **scribble**: Scribble to image
+        - **tile**: Upscaling and detail enhancement
+        
+        ## 💎 Popular LoRA Combinations
+        
+        ### For Portraits
+        - Base: `digiplay/majicMIX_realistic_v7`
+        - LoRA: `Detail Tweaker` or `Face Detail`
+        - VAE: `SD1.5 VAE`
+        
+        ### For Anime Characters
+        - Base: `digiplay/ChikMix_V3`
+        - LoRA: `Anime Art` or `Manga Style`
+        - VAE: `Anime VAE`
+        
+        ### For NSFW Content
+        - Base: Any NSFW-capable model
+        - LoRA: `NSFW Master`, `Realistic NSFW`, or `Anime NSFW`
+        - Note: Always use responsibly and legally
+        
+        ## ⚙️ Parameter Guidelines
+        
+        ### Steps
+        - **20-30**: Fast, good quality
+        - **30-50**: Balanced
+        - **50-100**: High quality, slow
+        
+        ### CFG Scale
+        - **5-7**: Creative, loose interpretation
+        - **7-10**: Balanced
+        - **10-15**: Strict prompt adherence
+        - **15+**: Very strict, may oversaturate
+        
+        ### Resolution
+        - **SD1.5**: 512x512, 512x768, 768x512
+        - **SDXL**: 1024x1024, 1024x1536, 1536x1024
+        
+        ## 🔞 NSFW Generation Guidelines
+        
+        1. Use NSFW-capable base models
+        2. Apply relevant LoRA for style enhancement
+        3. Use detailed prompts
+        4. Adjust CFG scale (7-12 recommended)
+        5. Consider using higher steps (40-60)
+        6. **Always comply with local laws and regulations**
+        
+        ## 🚀 Performance Tips
+        
+        - Unload models when switching between different types
+        - Use lower resolutions for testing
+        - Enable xFormers if available (automatic)
+        - Use appropriate batch sizes for your GPU
+        - Monitor GPU memory usage
         """)
-        
-        with gr.Row():
-            with gr.Column():
-                gr.Markdown("**SD1.5 Models Count:** " + str(len(SD15_MODELS)))
-                gr.Markdown("**SDXL Models Count:** " + str(len(SDXL_MODELS)))
-                gr.Markdown("**Chinese Models Count:** " + str(len(CHINESE_MODELS)))
-                gr.Markdown("**Florence-2 Models:** " + str(len(FLORENCE2_MODELS)))
-                gr.Markdown("**ControlNet Types:** " + str(len(CONTROLNET_MODELS) + len(SDXL_CONTROLNET_MODELS)))
-                gr.Markdown("**LoRA Models:** " + str(len(LORA_MODELS) - 1))  # Subtract "None"
-        
-        with gr.Row():
-            refresh_btn = gr.Button("🔄 Refresh Memory Info")
-            memory_info = gr.Textbox(label="Memory Status")
-        
-        def get_memory_info():
-            info = ""
-            if torch.cuda.is_available():
-                allocated = torch.cuda.memory_allocated() / 1024**3
-                reserved = torch.cuda.memory_reserved() / 1024**3
-                max_allocated = torch.cuda.max_memory_allocated() / 1024**3
-                info = f"Allocated: {allocated:.2f} GB\n"
-                info += f"Reserved: {reserved:.2f} GB\n"
-                info += f"Max Allocated: {max_allocated:.2f} GB"
-            else:
-                info = "Running on CPU - No GPU memory info"
-            return info
-        
-        refresh_btn.click(get_memory_info, outputs=memory_info)
 
 try:
     demo.launch(
-        server_name="0.0.0.0", 
-        server_port=7860, 
+        server_name="0.0.0.0",
+        server_port=7860,
         share=False,
         show_error=True,
         quiet=False