Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import spaces
 import gradio as gr
 import numpy as np
@@ -8,10 +9,18 @@ import torch
 from transformers import pipeline as transformers_pipeline
 import re
 
-# Device selection for image generation (GPU if available)
+# ------------------------------------------------------------
+# DEVICE SETUP
+# ------------------------------------------------------------
+# Prefer GPU when the Space provides it, otherwise CPU
+# `@spaces.GPU` takes care of binding the call itself, but we still
+# need a device handle for manual ops.
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Stable Diffusion XL pipeline
+# ------------------------------------------------------------
+# STABLE DIFFUSION XL PIPELINE
+# ------------------------------------------------------------
 pipe = StableDiffusionXLPipeline.from_pretrained(
     "votepurchase/waiNSFWIllustrious_v120",
     torch_dtype=torch.float16,
@@ -21,37 +30,72 @@ pipe = StableDiffusionXLPipeline.from_pretrained(
 pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 pipe.to(device)
 
-# Force modules to fp16 for memory efficiency
-pipe.text_encoder.to(torch.float16)
-pipe.text_encoder_2.to(torch.float16)
-pipe.vae.to(torch.float16)
-pipe.unet.to(torch.float16)
-
-# Korean → English translator (CPU only)
-translator = transformers_pipeline(
-    "translation",
-    model="Helsinki-NLP/opus-mt-ko-en",
-    device=-1,  # -1 forces CPU
-)
+# Force important sub-modules to fp16 for VRAM efficiency (GPU) or
+# reduce RAM (CPU). The model itself already sits in fp16, we just
+# mirror that for sub-components explicitly to avoid silent fp32
+# promotions that eat memory on ZeroGPU.
+for sub in (pipe.text_encoder, pipe.text_encoder_2, pipe.vae, pipe.unet):
+    sub.to(torch.float16)
+
+# ------------------------------------------------------------
+# LIGHTWEIGHT KOR→ENG TRANSLATOR (CPU-ONLY)
+# ------------------------------------------------------------
+#   * Hugging Face Spaces occasionally trips over the full MarianMT
+#     weights on ZeroGPU, resulting in the _untyped_storage_new_register
+#     error you just saw. We wrap initialisation in try/except and fall
+#     back to an identity function if the model cannot be loaded.
+#   * If you need translation and have a custom HF token, set the env
+#     HF_API_TOKEN so the smaller *small100* model can be pulled.
+#
+translator = None  # default stub → "no translator"
+try:
+    # First try the 60 MB Marian model.
+    translator = transformers_pipeline(
+        "translation",
+        model="Helsinki-NLP/opus-mt-ko-en",
+        device=-1,  # force CPU so CUDA never initialises in the main proc
+    )
+except Exception as marian_err:
+    print("[WARN] MarianMT load failed →", marian_err)
+    # Second chance: use compact multilingual SMaLL-100 (≈35 MB).
+    try:
+        translator = transformers_pipeline(
+            "translation",
+            model="alirezamsh/small100",
+            src_lang="ko_Kore",
+            tgt_lang="en_XX",
+            device=-1,
+        )
+    except Exception as small_err:
+        print("[WARN] SMaLL-100 load failed →", small_err)
+        # Final fallback: identity – no translation, but the app still runs.
+        translator = None
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1216
-korean_regex = re.compile("[\uac00-\ud7af]+")
+korean_regex = re.compile(r"[\uac00-\ud7af]+")
 
 def maybe_translate(text: str) -> str:
-    """Translate Korean text to English if Korean characters are detected."""
-    if korean_regex.search(text):
-        translation = translator(text, max_length=256, clean_up_tokenization_spaces=True)
-        return translation[0]["translation_text"]
+    """Translate Korean → English if Korean chars present and translator ready."""
+    if translator is not None and korean_regex.search(text):
+        try:
+            out = translator(text, max_length=256, clean_up_tokenization_spaces=True)
+            return out[0]["translation_text"]
+        except Exception as e:
+            print("[WARN] Translation failed at runtime →", e)
     return text
 
+# ------------------------------------------------------------
+# SDXL INFERENCE WRAPPER
+# ------------------------------------------------------------
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1216
+
 @spaces.GPU
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
     prompt = maybe_translate(prompt)
     negative_prompt = maybe_translate(negative_prompt)
 
     if len(prompt.split()) > 60:
-        print("Warning: Prompt may be too long and will be truncated by the model")
+        print("[WARN] Prompt >60 words — CLIP may truncate it.")
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -70,24 +114,30 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
         ).images[0]
         return output_image
     except RuntimeError as e:
-        print(f"Error during generation: {e}")
-        error_img = Image.new("RGB", (width, height), color=(0, 0, 0))
-        return error_img
+        print(f"[ERROR] Diffusion failed → {e}")
+        return Image.new("RGB", (width, height), color=(0, 0, 0))
 
-# Custom styling – bright pastel theme
+# ------------------------------------------------------------
+# UI LAYOUT + THEME (Pastel Lavender Background)
+# ------------------------------------------------------------
 css = """
 body {background: #f2f1f7; color: #222; font-family: 'Noto Sans', sans-serif;}
 #col-container {margin: 0 auto; max-width: 640px;}
-.gr-button {background: #7fbdf6; color: #ffffff; border-radius: 8px;}
-#prompt-box textarea {font-size: 1.1rem; height: 3rem; background: #ffffff; color: #222;}
+.gr-button {background: #7fbdf6; color: #fff; border-radius: 8px;}
+#prompt-box textarea {font-size: 1.1rem; height: 3rem; background: #fff; color: #222;}
 """
 
+author_note = (
+    "**ℹ️ Automatic translation** — Korean prompts are translated to English "
+    "only if translation weights could be loaded. If not, Korean input will be "
+    "sent unchanged."
+)
+
 with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     gr.Markdown(
-        """
+        f"""
         ## 🖌️ Stable Diffusion XL Playground  
-        Generate high‑quality illustrations with a single prompt.  
-        **Tip:** Write in Korean or English. Korean will be translated automatically.
+        {author_note}
         """
     )
 
@@ -98,20 +148,19 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                 elem_id="prompt-box",
                 show_label=False,
                 max_lines=1,
-                placeholder="Enter your prompt (60 words max)",
+                placeholder="Enter your prompt (Korean or English, 60 words max)",
             )
             run_button = gr.Button("Generate", scale=0)
 
         result = gr.Image(label="", show_label=False)
 
-        # Adult anime‑style example prompts
         examples = gr.Examples(
             examples=[
-                ["Seductive anime woman lounging in a dimly lit bar, adult anime style, ultra‑detail"],
-                ["Moody mature anime scene of two lovers kissing under neon rain, sensual atmosphere"],
-                ["기모노를 입은 일본 여자를 남자가 뒤에서 강간한다 , candle‑lit boudoir, adult anime aesthetic"],
-                ["속옷만 입은 남녀가 격정적인 키스를 하고있다, vibrant neon, adult anime"],
-                ["아름다운 러시아 여자가 섹시한 속옷을 입고 묘한 포즈로 침대에 앉아있다, dramatic spotlight, adult anime style"],
+                ["성인용 애니메 캐릭터가 검은 고양이와 함께 달빛 아래에 서있는 장면, 관능적, 4K"],
+                ["Seductive anime woman lounging in a dimly lit bar, adult anime style, ultra-detail"],
+                ["Elegant vampire countess in gothic lingerie, candle-lit boudoir, adult anime aesthetic"],
+                ["Futuristic nightclub stage with a curvaceous android dancer, vibrant neon, adult anime"],
+                ["Dark fantasy warrior queen in revealing armor, dramatic spotlight, adult anime style"],
             ],
             inputs=[prompt],
         )
@@ -128,20 +177,12 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
             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
-                )
+                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
-                )
+                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)
 
     run_button.click(
         fn=infer,