Update app.py

app.py

@@ -10,7 +10,6 @@ from PIL import Image
 from pathlib import Path
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 from huggingface_hub import hf_hub_download
-from transformers import CLIPVisionModelWithProjection
 from compel import Compel, ReturnedEmbeddingsType
 
 # --- Configuration ---
@@ -50,6 +49,11 @@ UPSCALER_URL = os.environ.get("UPSCALER_URL", "https://huggingface.co/spaces/har
 _char_cache = {}
 
 # --- Helper Functions ---
+def fix_vae_dtype(vae, dtype=torch.float16):
+    vae = vae.to(dtype=dtype)
+    vae.config.force_upcast = False
+    return vae
+
 def auto_res(input_image):
     if input_image is None: 
         return gr.update(value=1024), gr.update(value=1024), gr.update(value="No image — using default 1024 x 1024"), gr.update(value=None)
@@ -93,13 +97,6 @@ def encode_prompts(pipe, prompt, neg):
 # --- Preloading Logic ---
 def download_and_load_all():
     print("🚀 [Startup] Initializing Character Studio pre-load...")
-    
-    print("⚙️ [Pipeline] Loading ViT-H Image Encoder for IP-Adapter...")
-    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
-        "h94/IP-Adapter", 
-        subfolder="models/image_encoder", 
-        torch_dtype=torch.float16
-    )
 
     for model_name in CIVITAI_MODELS.keys():
         try:
@@ -116,19 +113,19 @@ def download_and_load_all():
                 )
             
             if model_name not in _char_cache:
-                print(f"⚙️ [Pipeline] Building {model_name} with IP-Adapter...")
+                print(f"⚙️ [Pipeline] Building {model_name} with Standard IP-Adapter...")
                 pipe = StableDiffusionXLPipeline.from_single_file(
                     str(dest), 
                     torch_dtype=torch.float16,
-                    use_safetensors=True,
-                    image_encoder=image_encoder 
+                    use_safetensors=True
                 )
                 pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
                 
+                # SWITCH TO STANDARD IP-ADAPTER: Much better for Anime styles
                 pipe.load_ip_adapter(
                     "h94/IP-Adapter",
                     subfolder="sdxl_models",
-                    weight_name="ip-adapter-plus_sdxl_vit-h.safetensors"
+                    weight_name="ip-adapter_sdxl.safetensors" 
                 )
                 _char_cache[model_name] = pipe
                 
@@ -140,6 +137,10 @@ def download_and_load_all():
 # --- Inference ---
 @spaces.GPU(duration=60)
 def generate_gpu(model_name, ref_img, prompt, neg, seed, randomize, w, h, ip_scale, cfg, steps, progress=gr.Progress(track_tqdm=False)):
+    if ref_img is None:
+        yield None, seed, gr.update(value="⚠️ Error: Please upload a Character Reference image first!"), gr.update(visible=False)
+        return
+
     if randomize: seed = random.randint(0, MAX_SEED)
     yield None, seed, gr.update(value="Waiting for GPU..."), gr.update(visible=False, value=None)
     
@@ -148,25 +149,34 @@ def generate_gpu(model_name, ref_img, prompt, neg, seed, randomize, w, h, ip_sca
         yield None, seed, gr.update(value=f"Model {model_name} not loaded."), gr.update(visible=False)
         return
 
-    # 1. FIX THE VAE UPCAST BUG PERMANENTLY
-    pipe.to("cuda")
-    pipe.vae.to(dtype=torch.float16, non_blocking=True)
+    try:
+        # 1. GPU Prep & VAE Lock (Fixes distorted colors)
+        pipe.to("cuda")
+        pipe.vae = fix_vae_dtype(pipe.vae)
+        
+        # 2. Disable the SDXL Watermarker (Prevents anime deep-frying)
+        pipe.watermark = None 
 
-    pipe.set_ip_adapter_scale(float(ip_scale))
-    generator = torch.Generator("cuda").manual_seed(seed)
-    
-    yield None, seed, gr.update(value=f"Generating {w}x{h}..."), gr.update(visible=False)
-    
-    # 2. AUTO-ENHANCE PROMPT FOR ILLUSTRIOUS/REALMIX
-    # This prevents the model from generating random noise when given short prompts like "running"
-    final_prompt = prompt
-    if "masterpiece" not in final_prompt.lower():
-        final_prompt = f"{prompt}, 1girl, solo, masterpiece, best quality, very aesthetic, absurdres"
+        pipe.set_ip_adapter_scale(float(ip_scale))
+        generator = torch.Generator("cuda").manual_seed(seed)
+        
+        # 3. Auto-enhance Prompt
+        final_prompt = prompt or ""
+        if "masterpiece" not in final_prompt.lower():
+            final_prompt = f"{final_prompt}, 1girl, solo, masterpiece, best quality, very aesthetic, absurdres"
 
-    def cb(p, i, t, kw): progress((i+1)/steps, desc=f"Step {i+1}/{steps}"); return kw
+        # 4. Auto-enhance Negative Prompt
+        base_neg = "lowres, (bad), text, error, missing, extra, worst quality, jpeg artifacts, bad anatomy, watermark, unfinished, displeasing, oldest, early, signature, extra digits, fewer digits, bad hands"
+        final_neg = f"{neg}, {base_neg}" if neg else base_neg
 
-    try:
-        cond, pool, ncond, npool = encode_prompts(pipe, final_prompt, neg)
+        yield None, seed, gr.update(value=f"Generating {w}x{h}..."), gr.update(visible=False)
+        
+        def cb(p, i, t, kw): 
+            progress((i+1)/steps, desc=f"Step {i+1}/{steps}")
+            return kw
+
+        cond, pool, ncond, npool = encode_prompts(pipe, final_prompt, final_neg)
+        
         image = pipe(
             prompt_embeds=cond,
             pooled_prompt_embeds=pool,
@@ -181,12 +191,14 @@ def generate_gpu(model_name, ref_img, prompt, neg, seed, randomize, w, h, ip_sca
             callback_on_step_end=cb
         ).images[0]
         
-    except RuntimeError as e:
-        torch.cuda.empty_cache()
-        yield None, seed, gr.update(value=f"Error: {str(e)[:120]}"), gr.update(visible=False)
-        return
+        yield image, seed, gr.update(value="✅ Done!"), gr.update(visible=True) 
         
-    yield image, seed, gr.update(value="Done!"), gr.update(visible=False)
+    except Exception as e:
+        import traceback
+        traceback.print_exc() 
+        torch.cuda.empty_cache()
+        yield None, seed, gr.update(value=f"❌ Crash: {str(e)[:150]}"), gr.update(visible=False)
+
 
 # --- UI Setup ---
 def build_css(accent="#7c3aed"):
@@ -229,7 +241,7 @@ with gr.Blocks(title="Character Lab", css=build_css(DEFAULT_ACCENT), theme=gr.th
             <a href="{UPSCALER_URL}" target="_blank">Upscaler</a>
         </div>
         ''')
-        gr.Markdown("# Character Lab\n**Maintain character consistency using IP-Adapter Plus.**")
+        gr.Markdown("# Character Lab\n**Maintain character consistency using Standard IP-Adapter.**")
 
         gr.HTML("<p class='lbl'>Select Model</p>")
         with gr.Row():
@@ -262,8 +274,8 @@ with gr.Blocks(title="Character Lab", css=build_css(DEFAULT_ACCENT), theme=gr.th
         with gr.Accordion("Advanced Settings", open=False):
             ch_neg = gr.Text(label="Negative Prompt", max_lines=2, value="lowres, bad anatomy, bad hands, text, watermark, worst quality")
             
-            # 3. LOWERED DEFAULT IP SCALE TO PREVENT DEEP-FRYING ANIME MODELS
-            ch_ip_scale = gr.Slider(label="Consistency Strength (IP-Scale)", minimum=0.1, maximum=1.2, step=0.05, value=0.55)
+            # Reset to a balanced default scale since Standard IP-Adapter is gentler
+            ch_ip_scale = gr.Slider(label="Consistency Strength (IP-Scale)", minimum=0.1, maximum=1.2, step=0.05, value=0.65)
             
             with gr.Row():
                 ch_seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)