Update app.py

app.py

@@ -11,14 +11,7 @@ import re
 from cohere import ClientV2
 
 # ------------------------------------------------------------
-# GLOBAL VARIABLES FOR LAZY LOADING
-# ------------------------------------------------------------
-pipe = None
-img2img_pipe = None
-device = None
-
-# ------------------------------------------------------------
-# INITIALIZE COHERE CLIENT FOR TRANSLATIONS AND PROMPT GENERATION
+# COHERE CLIENT
 # ------------------------------------------------------------
 coh_api_key = os.getenv("COH_API")
 if not coh_api_key:
@@ -34,84 +27,15 @@ else:
 
 
 # ------------------------------------------------------------
-# LAZY LOADING FUNCTIONS FOR PIPELINES
+# MODEL CONFIGURATION
 # ------------------------------------------------------------
-def load_pipeline():
-    """Lazily load the text-to-image pipeline on first use"""
-    global pipe, device
-    
-    if pipe is not None:
-        return pipe
-    
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    print(f"[INFO] Loading Text-to-Image pipeline on device: {device}")
-    
-    if torch.cuda.is_available():
-        pipe = StableDiffusionXLPipeline.from_pretrained(
-            "Heartsync/NSFW-Uncensored",
-            torch_dtype=torch.float16,
-            variant="fp16",
-            use_safetensors=True,
-        )
-        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-        pipe.to(device)
-        
-        # Force sub-modules to fp16 for VRAM efficiency
-        for sub in (pipe.text_encoder, pipe.text_encoder_2, pipe.vae, pipe.unet):
-            if sub is not None:
-                sub.to(torch.float16)
-    else:
-        pipe = StableDiffusionXLPipeline.from_pretrained(
-            "Heartsync/NSFW-Uncensored",
-            torch_dtype=torch.float32,
-            use_safetensors=True,
-        )
-        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-        pipe.to(device)
-    
-    print("[INFO] Text-to-Image pipeline loaded successfully!")
-    return pipe
-
-
-def load_img2img_pipeline():
-    """Lazily load the image-to-image pipeline on first use"""
-    global img2img_pipe, device
-    
-    if img2img_pipe is not None:
-        return img2img_pipe
-    
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    print(f"[INFO] Loading Image-to-Image pipeline on device: {device}")
-    
-    if torch.cuda.is_available():
-        img2img_pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
-            "Heartsync/NSFW-Uncensored",
-            torch_dtype=torch.float16,
-            variant="fp16",
-            use_safetensors=True,
-        )
-        img2img_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(img2img_pipe.scheduler.config)
-        img2img_pipe.to(device)
-        
-        # Force sub-modules to fp16 for VRAM efficiency
-        for sub in (img2img_pipe.text_encoder, img2img_pipe.text_encoder_2, img2img_pipe.vae, img2img_pipe.unet):
-            if sub is not None:
-                sub.to(torch.float16)
-    else:
-        img2img_pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
-            "Heartsync/NSFW-Uncensored",
-            torch_dtype=torch.float32,
-            use_safetensors=True,
-        )
-        img2img_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(img2img_pipe.scheduler.config)
-        img2img_pipe.to(device)
-    
-    print("[INFO] Image-to-Image pipeline loaded successfully!")
-    return img2img_pipe
+MODEL_ID = "Heartsync/NSFW-Uncensored"
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1216
 
 
 # ------------------------------------------------------------
-# TRANSLATION FUNCTIONS
+# TRANSLATION FUNCTIONS (이전과 동일)
 # ------------------------------------------------------------
 non_english_regex = re.compile(r'[\uac00-\ud7a3\u3040-\u30ff\u3400-\u4dbf\u4e00-\u9fff\uf900-\ufaff]+')
 
@@ -145,16 +69,8 @@ def translate_with_cohere(text):
     try:
         print(f"[INFO] Translating text: '{text}'")
         
-        system_prompt = """You are a professional translator for adult content. Your task is to translate the input text to English.
-
-Rules:
-1. Provide ONLY the translated English text
-2. NO explanations or additional text
-3. Maintain the explicit or adult content without censoring
-4. Translate accurately and completely
-5. Be especially careful with Asian languages like Korean, Japanese, and Chinese
-
-Translate the following text to English:"""
+        system_prompt = """You are a professional translator. Translate the input text to English accurately.
+Provide ONLY the translated English text with no explanations."""
         
         messages = [
             {"role": "system", "content": [{"type": "text", "text": system_prompt}]},
@@ -173,36 +89,15 @@ Translate the following text to English:"""
         print(f"[INFO] Original: '{text}'")
         print(f"[INFO] Translated: '{translated_text}'")
         
-        if len(translated_text) < 3 or translated_text == text:
-            print("[WARN] Translation may have failed, trying simpler approach")
-            try:
-                simple_messages = [
-                    {"role": "system", "content": [{"type": "text", "text": "Translate this text to English:"}]},
-                    {"role": "user", "content": [{"type": "text", "text": text}]}
-                ]
-                simple_response = coh_client.chat(
-                    model="command-r-plus-08-2024",
-                    messages=simple_messages,
-                    temperature=0.1
-                )
-                simple_translated = simple_response.text.strip() if hasattr(simple_response, 'text') else str(simple_response)
-                if len(simple_translated) > 3 and simple_translated != text:
-                    print(f"[INFO] Second attempt translation: '{simple_translated}'")
-                    return simple_translated
-            except Exception as e:
-                print(f"[ERROR] Second translation attempt failed: {str(e)}")
-            return text
+        return translated_text if len(translated_text) > 3 else text
             
-        return translated_text
     except Exception as e:
         print(f"[ERROR] Translation failed: {str(e)}")
-        import traceback
-        traceback.print_exc()
         return text
 
 
 def translate_prompt_if_needed(prompt):
-    """Helper function to translate prompt using Cohere"""
+    """Helper function to translate prompt"""
     if not is_non_english(prompt):
         return prompt
     
@@ -210,7 +105,7 @@ def translate_prompt_if_needed(prompt):
         return prompt
     
     try:
-        trans_system = "You are a translator. Translate the following text to English accurately. Only provide the translation, no comments or explanations."
+        trans_system = "Translate to English accurately. Only provide the translation."
         
         trans_response = coh_client.chat(
             model="command-r-plus-08-2024",
@@ -221,22 +116,8 @@ def translate_prompt_if_needed(prompt):
             temperature=0.1
         )
         
-        translated_prompt = None
-        
         if hasattr(trans_response, 'text'):
-            translated_prompt = trans_response.text
-        elif hasattr(trans_response, 'response'):
-            translated_prompt = trans_response.response
-        elif isinstance(trans_response, dict) and 'text' in trans_response:
-            translated_prompt = trans_response['text']
-        else:
-            response_str = str(trans_response)
-            match = re.search(r"text=['\"](.*?)['\"]", response_str)
-            if match:
-                translated_prompt = match.group(1)
-        
-        if translated_prompt:
-            translated_prompt = translated_prompt.strip()
+            translated_prompt = trans_response.text.strip()
             print(f"[SUCCESS] Translated: '{prompt}' -> '{translated_prompt}'")
             return translated_prompt
             
@@ -251,35 +132,11 @@ def translate_prompt_if_needed(prompt):
 # ------------------------------------------------------------
 prompt_examples = [
     "The shy college girl, with glasses and a tight plaid skirt, nervously approaches her professor",
-    "Her skirt rose a little higher with each gentle push, a soft blush of blush spreading across her cheeks as she felt the satisfying warmth of his breath on her cheek.",
-    "a girl in a school uniform having her skirt pulled up by a boy, and then being fucked",
-    "Moody mature anime scene of two lovers fuck under neon rain, sensual atmosphere",
-    "Moody mature anime scene of two lovers kissing under neon rain, sensual atmosphere",
-    "The girl sits on the boy's lap by the window, his hands resting on her waist. She is unbuttoning his shirt, her expression focused and intense.",
-    "A girl with long, black hair is sleeping on her desk in the classroom. Her skirt has ridden up, revealing her thighs, and a trail of drool escapes her slightly parted lips.",
-    "The waves rolled gently, a slow, sweet kiss of the lip, a slow, slow build of anticipation as their toes bumped gently – a slow, sweet kiss of the lip, a promise of more to come.",
-    "Her elegant silk gown swayed gracefully as she approached him, the delicate fabric brushing against her legs. A warm blush spread across her cheeks as she felt his breath on her face.",
-    "Her white blouse and light cotton skirt rose a little higher with each gentle push, a soft blush spreading across her cheeks as she felt the satisfying warmth of his breath on her cheek.",
-    "A woman in a business suit having her skirt lifted by a man, and then being sexually assaulted.",
-    "The older woman sits on the man's lap by the fireplace, his hands resting on her hips. She is unbuttoning his vest, her expression focused and intense. He takes control of the situation as she finishes unbuttoning his shirt, pushing her onto her back and begins to have sex with her.",
-    "There is a woman with long black hair. Her face features alluring eyes and full lips, with a slender figure adorned in black lace lingerie. She lies on the bed, loosening her lingerie strap with one hand while seductively glancing downward.",
-    "In a dimly lit room, the same woman teases with her dark, flowing hair, now covering her voluptuous breasts, while a black garter belt accentuates her thighs. She sits on the sofa, leaning back, lifting one leg to expose her most private areas through the sheer lingerie.",
-    "A woman with glasses, lying on the bed in just her bra, spreads her legs wide, revealing all! She wears a sultry expression, gazing directly at the viewer with her brown eyes, her short black hair cascading over the pillow. Her slim figure, accentuated by the lacy lingerie, exudes a seductive aura.",
-    "A soft focus on the girl's face, eyes closed, biting her lip, as her roommate performs oral pleasure, the experienced woman's hair cascading between her thighs.",
-    "A woman in a blue hanbok sits on a wooden floor, her legs folded beneath her, gazing out of a window, the sunlight highlighting the graceful lines of her clothing.",
-    "The couple, immersed in a wooden outdoor bath, share an intimate moment, her wet kimono clinging to her curves, his hands exploring her body beneath the water's surface.",
-    "A steamy shower scene, the twins embrace under the warm water, their soapy hands gliding over each other's curves, their passion intensifying as they explore uncharted territories.",
-    "The teacher, with a firm grip, pins the student against the blackboard, her skirt hiked up, exposing her delicate lace panties. Their heavy breathing echoes in the quiet room as they share an intense, intimate moment.",
-    "After hours, the girl sits on top of the teacher's lap, riding him on the classroom floor, her hair cascading over her face as she moves with increasing intensity, their bodies glistening with sweat.",
-    "In the dimly lit dorm room, the roommates lay entangled in a passionate embrace, their naked bodies glistening with sweat, as the experienced woman teaches her lover the art of kissing and touching.",
-    "The once-innocent student, now confident, takes charge, straddling her lover on the couch, their bare skin illuminated by the warm glow of the sunset through the window.",
-    "A close-up of the secretary's hand unzipping her boss's dress shirt, her fingers gently caressing his chest, their eyes locked in a heated embrace in the supply closet.",
-    "The secretary, in a tight pencil skirt and silk blouse, leans back on the boss's desk, her legs wrapped around his waist, her blouse unbuttoned, revealing her lace bra, as he passionately kisses her, his hands exploring her body.",
-    "On the living room couch, one twin sits astride her sister's lap, their lips locked in a passionate kiss, their hands tangled in each other's hair, unraveling a new level of intimacy.",
-    "In a dimly lit chamber, the dominant woman, dressed in a leather corset and thigh-high boots, stands tall, her hand gripping her submissive partner's hair, his eyes closed in submission as she instructs him to please her.",
-    "The dominant, in a sheer lace bodysuit, sits on a throne-like chair, her legs spread, as the submissive, on his knees, worships her with his tongue, his hands bound behind his back.",
-    "A traditional Japanese onsen, with steam rising, a young woman in a colorful kimono kneels on a tatami mat, her back to the viewer, as her male partner, also in a kimono, gently unties her obi, revealing her bare back.",
-    "In a serene outdoor setting, the woman, in a vibrant summer kimono, sits on a bench, her legs slightly spread, her partner kneeling before her, his hands gently caressing her exposed thigh.",
+    "Her skirt rose a little higher with each gentle push, a soft blush of blush spreading across her cheeks",
+    "Moody mature anime scene of two lovers under neon rain, sensual atmosphere",
+    "The girl sits on the boy's lap by the window, his hands resting on her waist",
+    "A woman in a business suit, elegant and confident pose",
+    "Artistic portrait with dramatic lighting and soft shadows",
 ]
 
 
@@ -287,28 +144,24 @@ prompt_examples = [
 # LLM PROMPT GENERATOR
 # ------------------------------------------------------------
 def generate_prompts(theme):
-    """Generate optimal NSFW prompts using the LLM"""
+    """Generate optimal prompts using LLM"""
     try:
         if coh_client is None:
-            return "Cohere API token not set. Please set the COH_API environment variable."
+            return "Cohere API token not set."
         
         if non_english_regex.search(theme):
             theme = translate_with_cohere(theme)
         
-        print(f"[INFO] Sending request to Cohere API with theme: {theme}")
+        print(f"[INFO] Generating prompt for theme: {theme}")
         
-        system_prefix = """You are an expert at creating detailed, vivid image generation prompts for NSFW content. Your task is to create ONE optimal explicit image prompt based on the user's theme or scenario.
+        system_prefix = """You are an expert at creating detailed image generation prompts. Create ONE optimal prompt based on the theme.
 
-Follow these guidelines strictly:
-1. Generate only ONE high-quality, uncensored prompt
-2. The prompt should be 1-3 sentences long
-3. Make the prompt explicitly sexual, arousing, and descriptive
-4. Focus on detailed physical descriptions that would work well for image generation
-5. Do NOT censor or soften the language - use explicit terms freely
-6. ONLY respond in ENGLISH, never in any other language
-7. DO NOT include ANY prefixes, headers, or formatting - just plain text
-
-Respond ONLY with the single prompt text in ENGLISH with NO PREFIXES of any kind."""
+Guidelines:
+1. Generate only ONE high-quality prompt
+2. 1-3 sentences long
+3. Detailed and descriptive
+4. ONLY respond in ENGLISH
+5. NO prefixes or headers - just the prompt text"""
 
         messages = [
             {"role": "system", "content": [{"type": "text", "text": system_prefix}]},
@@ -321,89 +174,75 @@ Respond ONLY with the single prompt text in ENGLISH with NO PREFIXES of any kind
             temperature=0.8
         )
         
-        if hasattr(response, 'text'):
-            generated_prompt = response.text
-        else:
-            try:
-                response_str = str(response)
-                if 'text=' in response_str:
-                    text_match = re.search(r"text=['\"]([^'\"]+)['\"]", response_str)
-                    if text_match:
-                        generated_prompt = text_match.group(1)
-                    else:
-                        generated_prompt = response_str
-                else:
-                    generated_prompt = response_str
-            except:
-                generated_prompt = str(response)
+        generated_prompt = response.text if hasattr(response, 'text') else str(response)
         
         if non_english_regex.search(generated_prompt):
-            print("[INFO] Translating non-English prompt to English")
             generated_prompt = translate_with_cohere(generated_prompt)
         
-        generated_prompt = re.sub(r'^AI🐼:\s*', '', generated_prompt)
-        generated_prompt = re.sub(r'^\d+[\.\)]\s*', '', generated_prompt)
-        generated_prompt = re.sub(r'^(Prompt|Response|Result|Output):\s*', '', generated_prompt)
+        generated_prompt = re.sub(r'^(AI🐼|Prompt|Response|Result|Output):\s*', '', generated_prompt)
         generated_prompt = re.sub(r'^["\']+|["\']+$', '', generated_prompt)
         generated_prompt = generated_prompt.strip()
-        generated_prompt = re.sub(r'\s+', ' ', generated_prompt)
         
-        print(f"[INFO] Generated prompt: {generated_prompt}")
+        print(f"[INFO] Generated: {generated_prompt}")
         
-        if len(generated_prompt) > 10:
-            return generated_prompt
-        else:
-            return "Failed to generate a valid prompt"
+        return generated_prompt if len(generated_prompt) > 10 else "Failed to generate prompt"
         
     except Exception as e:
         print(f"[ERROR] Prompt generation failed: {str(e)}")
-        import traceback
-        traceback.print_exc()
-        return f"Error generating prompt: {str(e)}"
+        return f"Error: {str(e)}"
 
 
 # ------------------------------------------------------------
-# SDXL INFERENCE WRAPPER (Text-to-Image)
+# SDXL TEXT-TO-IMAGE (ZeroGPU 최적화)
 # ------------------------------------------------------------
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1216
-
-
 @spaces.GPU(duration=120)
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
-    """Text-to-Image generation with ZeroGPU support"""
-    global pipe, device
+    """Text-to-Image generation optimized for ZeroGPU"""
     
-    # Load pipeline inside GPU context
-    pipe = load_pipeline()
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"[DEBUG] Original prompt: '{prompt}'")
     
-    print(f"[DEBUG] Original prompt received: '{prompt}'")
-    print(f"[DEBUG] Device: {device}")
-    
-    # Translate prompts if needed
+    # Translate if needed
     if is_non_english(prompt):
-        print(f"[ALERT] Non-English prompt detected: '{prompt}'")
+        print(f"[ALERT] Non-English prompt detected")
         prompt = translate_prompt_if_needed(prompt)
-        print(f"[INFO] Final prompt: '{prompt}'")
+        print(f"[INFO] Translated prompt: '{prompt}'")
     
     if is_non_english(negative_prompt):
-        print(f"[ALERT] Non-English negative prompt detected: '{negative_prompt}'")
         negative_prompt = translate_prompt_if_needed(negative_prompt)
-        print(f"[INFO] Final negative prompt: '{negative_prompt}'")
     
-    print(f"[INFO] Final prompt to use: '{prompt}'")
-    print(f"[INFO] Final negative prompt to use: '{negative_prompt}'")
-
-    if len(prompt.split()) > 60:
-        print("[WARN] Prompt >60 words — CLIP may truncate it.")
-
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-
-    generator = torch.Generator(device=device).manual_seed(seed)
-
+    
     try:
+        # GPU 컨텍스트 내에서 파이프라인 로드
+        device = torch.device("cuda")
+        print(f"[INFO] Loading pipeline on device: {device}")
+        
+        # FP16 파이프라인 로드
+        pipe = StableDiffusionXLPipeline.from_pretrained(
+            MODEL_ID,
+            torch_dtype=torch.float16,
+            variant="fp16",
+            use_safetensors=True,
+        )
+        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+        pipe = pipe.to(device)
+        
+        # 메모리 효율성 개선
+        try:
+            pipe.enable_xformers_memory_efficient_attention()
+            print("[INFO] xformers memory efficient attention enabled")
+        except:
+            print("[WARN] xformers not available, using default attention")
+        
+        # VAE 타일링 활성화 (메모리 절약)
+        pipe.enable_vae_tiling()
+        
+        # 생성기 설정
+        generator = torch.Generator(device=device).manual_seed(seed)
+        
+        # 이미지 생성
+        print(f"[INFO] Generating image...")
         output_image = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -413,53 +252,77 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
             height=height,
             generator=generator,
         ).images[0]
+        
+        # 메모리 정리
+        del pipe
+        torch.cuda.empty_cache()
+        
+        print("[SUCCESS] Image generated successfully")
         return output_image, seed
-    except RuntimeError as e:
-        print(f"[ERROR] Diffusion failed → {e}")
+        
+    except Exception as e:
+        print(f"[ERROR] Generation failed: {str(e)}")
         import traceback
         traceback.print_exc()
+        
+        # 에러 발생 시 메모리 정리
+        torch.cuda.empty_cache()
+        
         return Image.new("RGB", (width, height), color=(0, 0, 0)), seed
 
 
 # ------------------------------------------------------------
-# SDXL INFERENCE WRAPPER (Image-to-Image)
+# SDXL IMAGE-TO-IMAGE (ZeroGPU 최적화)
 # ------------------------------------------------------------
 @spaces.GPU(duration=120)
 def img2img_infer(init_image, prompt, negative_prompt, strength, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
-    """Image-to-Image generation with ZeroGPU support"""
-    global img2img_pipe, device
+    """Image-to-Image generation optimized for ZeroGPU"""
     
     if init_image is None:
         return None, seed
     
-    # Load pipeline inside GPU context
-    img2img_pipe = load_img2img_pipeline()
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    
-    print(f"[DEBUG] Image-to-Image prompt received: '{prompt}'")
-    print(f"[DEBUG] Device: {device}")
+    print(f"[DEBUG] Image-to-Image prompt: '{prompt}'")
     
-    # Translate prompts if needed
+    # Translate if needed
     if is_non_english(prompt):
-        print(f"[ALERT] Non-English prompt detected: '{prompt}'")
         prompt = translate_prompt_if_needed(prompt)
-        print(f"[INFO] Translated prompt: '{prompt}'")
     
     if is_non_english(negative_prompt):
-        print(f"[ALERT] Non-English negative prompt detected: '{negative_prompt}'")
         negative_prompt = translate_prompt_if_needed(negative_prompt)
-        print(f"[INFO] Translated negative prompt: '{negative_prompt}'")
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     
-    generator = torch.Generator(device=device).manual_seed(seed)
-    
-    # Preprocess image
-    init_image = init_image.convert("RGB")
-    init_image = init_image.resize((width, height), Image.Resampling.LANCZOS)
-    
     try:
+        # GPU 컨텍스트 내에서 파이프라인 로드
+        device = torch.device("cuda")
+        print(f"[INFO] Loading img2img pipeline on device: {device}")
+        
+        img2img_pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+            MODEL_ID,
+            torch_dtype=torch.float16,
+            variant="fp16",
+            use_safetensors=True,
+        )
+        img2img_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(img2img_pipe.scheduler.config)
+        img2img_pipe = img2img_pipe.to(device)
+        
+        # 메모리 효율성 개선
+        try:
+            img2img_pipe.enable_xformers_memory_efficient_attention()
+        except:
+            pass
+        
+        img2img_pipe.enable_vae_tiling()
+        
+        # 이미지 전처리
+        init_image = init_image.convert("RGB")
+        init_image = init_image.resize((width, height), Image.Resampling.LANCZOS)
+        
+        generator = torch.Generator(device=device).manual_seed(seed)
+        
+        # 이미지 생성
+        print(f"[INFO] Transforming image...")
         output_image = img2img_pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -469,11 +332,21 @@ def img2img_infer(init_image, prompt, negative_prompt, strength, seed, randomize
             num_inference_steps=num_inference_steps,
             generator=generator,
         ).images[0]
+        
+        # 메모리 정리
+        del img2img_pipe
+        torch.cuda.empty_cache()
+        
+        print("[SUCCESS] Image transformed successfully")
         return output_image, seed
-    except RuntimeError as e:
-        print(f"[ERROR] Image-to-Image generation failed → {e}")
+        
+    except Exception as e:
+        print(f"[ERROR] Transformation failed: {str(e)}")
         import traceback
         traceback.print_exc()
+        
+        torch.cuda.empty_cache()
+        
         return None, seed
 
 
@@ -489,119 +362,45 @@ def boost_prompt(keyword):
         return "Please enter a keyword or theme first"
     
     if coh_client is None:
-        return "Cohere API token not set. Please set the COH_API environment variable."
+        return "Cohere API token not set"
     
-    print(f"[INFO] Generating boosted prompt for keyword: {keyword}")
     prompt = generate_prompts(keyword)
-    
-    if isinstance(prompt, str) and len(prompt) > 10 and not prompt.startswith("Error") and not prompt.startswith("Failed"):
-        return prompt.strip()
-    else:
-        return "Failed to generate a suitable prompt. Please try again with a different keyword."
+    return prompt.strip() if len(prompt) > 10 else "Failed to generate prompt"
 
 
 # ------------------------------------------------------------
-# UI LAYOUT + THEME
+# UI LAYOUT
 # ------------------------------------------------------------
 css = """
 body {background: linear-gradient(135deg, #f2e6ff 0%, #e6f0ff 100%); color: #222; font-family: 'Noto Sans', sans-serif;}
 #col-container {margin: 0 auto; max-width: 768px; padding: 15px; background: rgba(255, 255, 255, 0.8); border-radius: 15px; box-shadow: 0 8px 32px rgba(31, 38, 135, 0.2);}
 .gr-button {background: #7fbdf6; color: #fff; border-radius: 8px; transition: all 0.3s ease; font-weight: bold;}
 .gr-button:hover {background: #5a9ae6; transform: translateY(-2px); box-shadow: 0 5px 15px rgba(0,0,0,0.1);}
-#prompt-box textarea {font-size: 1.1rem; height: 9rem !important; background: #fff; color: #222; border-radius: 10px; border: 1px solid #d1c1e0;}
+#prompt-box textarea {font-size: 1.1rem; height: 9rem !important; background: #fff; color: #222; border-radius: 10px;}
 .boost-btn {background: #ff7eb6; margin-top: 5px;}
 .boost-btn:hover {background: #ff5aa5;}
 .random-btn {background: #9966ff; margin-top: 5px;}
 .random-btn:hover {background: #8040ff;}
-.container {animation: fadeIn 0.5s ease-in-out;}
 .title {color: #6600cc; text-shadow: 1px 1px 2px rgba(0,0,0,0.1);}
-.gr-form {border: none !important; background: transparent !important;}
-.gr-input {border-radius: 8px !important;}
-.gr-slider {height: 12px !important;}
-.gr-slider .handle {height: 20px !important; width: 20px !important;}
-.panel {border-radius: 12px; overflow: hidden; box-shadow: 0 4px 15px rgba(0,0,0,0.1);}
-.gr-image {border-radius: 12px; overflow: hidden; transition: all 0.3s ease;}
-.gr-image:hover {transform: scale(1.02); box-shadow: 0 8px 25px rgba(0,0,0,0.15);}
-@keyframes fadeIn {
-  from {opacity: 0; transform: translateY(20px);}
-  to {opacity: 1; transform: translateY(0);}
-}
-.gr-accordion {border-radius: 10px; overflow: hidden; transition: all 0.3s ease;}
-.gr-accordion:hover {box-shadow: 0 5px 15px rgba(0,0,0,0.1);}
 """
 
-author_note = "**ℹ️ A research platform pushing the limits of uncensored AI image generation. Input prompts in any language - automatic translation and generation supported.**"
-
-
-with gr.Blocks(
-    css=css, 
-    theme=gr.themes.Soft(),
-    head="""
-    <!-- Google tag (gtag.js) -->
-    <script async src="https://www.googletagmanager.com/gtag/js?id=G-GTFK201G22"></script>
-    <script>
-      window.dataLayer = window.dataLayer || [];
-      function gtag(){dataLayer.push(arguments);}
-      gtag('js', new Date());
-      gtag('config', 'G-GTFK201G22');
-    </script>
-    """
-) as demo:
+with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     gr.Markdown(
-        f"""
+        """
         ## 🖌️ NSFW Uncensored Text & Imagery: AI Limits Explorer  
         
-        **New Update: Image-to-Image functionality has been added as a new tab! Upload your images and experiment with various transformations.**
-        
-        {author_note}
+        **ZeroGPU Optimized | Multi-language Support**
         """, elem_classes=["title"]
     )
     
-    with gr.Group(elem_classes="model-description"):
-        gr.HTML("""
-        <p>
-        <strong>Models Use cases: </strong><br>
-        </p>
-        <div style="display: flex; justify-content: center; align-items: center; gap: 10px; flex-wrap: wrap; margin-top: 10px; margin-bottom: 20px;">
-            <a href="https://huggingface.co/spaces/Heartsync/FREE-NSFW-HUB" target="_blank">
-                <img src="https://img.shields.io/static/v1?label=huggingface&message=FREE%20NSFW%20HUB&color=%230000ff&labelColor=%23800080&logo=huggingface&logoColor=%23ffa500&style=for-the-badge" alt="badge">
-            </a>
-            <a href="https://huggingface.co/spaces/Heartsync/PornHUB" target="_blank">
-                <img src="https://img.shields.io/static/v1?label=Porn%20HUB&message=NSFW%20Uncensored&color=%23ffc0cb&labelColor=%23ffff00&logo=huggingface&logoColor=%23ffa500&style=for-the-badge" alt="badge">
-            </a>
-            <a href="https://huggingface.co/spaces/Heartsync/adult" target="_blank">
-                <img src="https://img.shields.io/static/v1?label=Text%20to%20Image%20to%20Video&message=ADULT&color=%23ff00ff&labelColor=%23000080&logo=Huggingface&logoColor=%23ffa500&style=for-the-badge" alt="badge">
-            </a>
-            <a href="https://www.humangen.ai" target="_blank">
-              <img src="https://img.shields.io/static/v1?label=100% FREE&message=AI%20Playground&color=%230000ff&labelColor=%23800080&logo=huggingface&logoColor=%23ffa500&style=for-the-badge" alt="badge">
-            </a>   
-            <a href="https://huggingface.co/spaces/Heartsync/NSFW-Uncensored-image" target="_blank">
-                <img src="https://img.shields.io/static/v1?label=Image%20to%20Video&message=NSFW%20Uncensored&color=%230000ff&labelColor=%23800080&logo=Huggingface&logoColor=%23ffa500&style=for-the-badge" alt="badge">
-            </a>            
-            <a href="https://huggingface.co/spaces/Heartsync/NSFW-Uncensored-video2" target="_blank">
-                <img src="https://img.shields.io/static/v1?label=Image%20to%20Video(Mirror)&message=NSFW%20Uncensored&color=%230000ff&labelColor=%23800080&logo=Huggingface&logoColor=%23ffa500&style=for-the-badge" alt="badge">
-            </a>
-        </div>
-        <p>
-        <small style="opacity: 0.8;">High-quality image generation powered by StableDiffusionXL with video generation capability. Supports long prompts and various artistic styles.</small>
-        </p>
-        """)
-    
-    # State variables
-    current_image = gr.State(None)
-    current_seed = gr.State(0)
-
-    # Tabs
     with gr.Tabs():
         # Text-to-Image Tab
         with gr.TabItem("Text to Image"):
-            with gr.Column(elem_id="col-container", elem_classes=["container", "panel"]):
+            with gr.Column(elem_id="col-container"):
                 with gr.Row():
                     keyword_input = gr.Text(
                         label="Keyword Input",
-                        show_label=True,
-                        max_lines=1,
-                        placeholder="Enter a keyword or theme in any language to generate an optimal prompt",
+                        placeholder="Enter keyword in any language",
                         value="random",
                     )
                     boost_button = gr.Button("BOOST", elem_classes=["boost-btn"])
@@ -611,131 +410,70 @@ with gr.Blocks(
                     prompt = gr.Text(
                         label="Prompt",
                         elem_id="prompt-box",
-                        show_label=True,
                         max_lines=3,
-                        placeholder="Enter your prompt in any language (Korean, English, Japanese, etc.)",
+                        placeholder="Enter prompt in any language",
                     )
                     run_button = gr.Button("Generate", scale=0)
 
-                result = gr.Image(label="Generated Image", elem_classes=["gr-image"])
+                result = gr.Image(label="Generated Image")
 
-                with gr.Accordion("Advanced Settings", open=False, elem_classes=["gr-accordion"]):
+                with gr.Accordion("Advanced Settings", open=False):
                     negative_prompt = gr.Text(
                         label="Negative prompt",
-                        max_lines=1,
-                        placeholder="Enter a negative prompt in any language",
-                        value="text, talk bubble, low quality, watermark, signature",
+                        value="low quality, watermark, signature",
                     )
-
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     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=1024)
-
                     with gr.Row():
                         guidance_scale = gr.Slider(label="Guidance scale", minimum=0.0, maximum=20.0, step=0.1, value=7)
-                        num_inference_steps = gr.Slider(label="Inference steps", minimum=1, maximum=28, step=1, value=28)
+                        num_inference_steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28)
 
         # Image-to-Image Tab
         with gr.TabItem("Image to Image"):
-            with gr.Column(elem_id="col-container", elem_classes=["container", "panel"]):
-                input_image = gr.Image(
-                    label="Input Image",
-                    type="pil",
-                    elem_classes=["gr-image"]
-                )
+            with gr.Column(elem_id="col-container"):
+                input_image = gr.Image(label="Input Image", type="pil")
                 
                 with gr.Row():
                     img2img_prompt = gr.Text(
                         label="Prompt",
-                        show_label=True,
-                        max_lines=3,
-                        placeholder="Describe how you want to transform the image (any language)",
+                        placeholder="Describe transformation (any language)",
                     )
                     img2img_run_button = gr.Button("Transform", scale=0)
                 
-                img2img_result = gr.Image(label="Transformed Image", elem_classes=["gr-image"])
+                img2img_result = gr.Image(label="Transformed Image")
                 
-                with gr.Accordion("Advanced Settings", open=False, elem_classes=["gr-accordion"]):
+                with gr.Accordion("Advanced Settings", open=False):
                     img2img_negative_prompt = gr.Text(
                         label="Negative prompt",
-                        max_lines=1,
-                        placeholder="What to avoid in the transformation",
-                        value="low quality, watermark, signature",
+                        value="low quality, watermark",
                     )
-                    
-                    strength = gr.Slider(
-                        label="Transformation Strength",
-                        minimum=0.0,
-                        maximum=1.0,
-                        step=0.01,
-                        value=0.75,
-                        info="Lower values preserve more of the original image"
-                    )
-                    
+                    strength = gr.Slider(label="Strength", minimum=0.0, maximum=1.0, step=0.01, value=0.75)
                     img2img_seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     img2img_randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-                    
                     with gr.Row():
                         img2img_width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
                         img2img_height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
-                    
                     with gr.Row():
-                        img2img_guidance_scale = gr.Slider(label="Guidance scale", minimum=0.0, maximum=20.0, step=0.1, value=7.5)
-                        img2img_num_inference_steps = gr.Slider(label="Inference steps", minimum=1, maximum=50, step=1, value=30)
-
-    # Helper function
-    def update_image_state(img, seed_val):
-        return img, seed_val
+                        img2img_guidance_scale = gr.Slider(label="Guidance", minimum=0.0, maximum=20.0, step=0.1, value=7.5)
+                        img2img_num_inference_steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=30)
 
     # Event handlers
-    boost_button.click(
-        fn=boost_prompt,
-        inputs=[keyword_input],
-        outputs=[prompt]
-    )
+    boost_button.click(fn=boost_prompt, inputs=[keyword_input], outputs=[prompt])
+    random_button.click(fn=get_random_prompt, outputs=[prompt])
     
-    random_button.click(
-        fn=get_random_prompt,
-        inputs=[],
-        outputs=[prompt]
-    )
-
     run_button.click(
         fn=infer,
-        inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, current_seed]
-    ).then(
-        fn=update_image_state,
-        inputs=[result, current_seed],
-        outputs=[current_image, current_seed]
+        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        outputs=[result, seed]
     )
-
+    
     img2img_run_button.click(
         fn=img2img_infer,
-        inputs=[
-            input_image,
-            img2img_prompt,
-            img2img_negative_prompt,
-            strength,
-            img2img_seed,
-            img2img_randomize_seed,
-            img2img_width,
-            img2img_height,
-            img2img_guidance_scale,
-            img2img_num_inference_steps
-        ],
+        inputs=[input_image, img2img_prompt, img2img_negative_prompt, strength, img2img_seed, 
+                img2img_randomize_seed, img2img_width, img2img_height, img2img_guidance_scale, img2img_num_inference_steps],
         outputs=[img2img_result, img2img_seed]
     )