Update app.py

app.py

@@ -1,3 +1,12 @@
+# ============================================================
+# Hugging Face Spaces GPU app
+# IMPORTANT:
+# - spaces MUST be imported first
+# - @spaces.GPU MUST be used directly
+# ============================================================
+
+import spaces  # MUST be first, no try/except
+
 import os
 import random
 import gc
@@ -6,13 +15,6 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 
-try:
-    import spaces
-    GPU_DECORATOR = spaces.GPU
-except Exception:
-    def GPU_DECORATOR(fn):
-        return fn
-
 import torch
 from diffusers import (
     StableDiffusionPipeline,
@@ -32,83 +34,61 @@ HF_TOKEN = os.getenv("HF_TOKEN", "").strip()
 if HF_TOKEN:
     login(token=HF_TOKEN)
 
-cuda_available = torch.cuda.is_available()
-device = torch.device("cuda" if cuda_available else "cpu")
-dtype = torch.float16 if cuda_available else torch.float32
+device = torch.device("cuda")
+dtype = torch.float16
 
 MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 768 if not cuda_available else 1024
+MAX_IMAGE_SIZE = 1024
 
 pipe_txt2img = None
 pipe_img2img = None
-model_loaded = False
-load_error = None
 
 # ============================================================
-# Load model (FORCED tokenizer fix)
+# Load model (once at startup)
 # ============================================================
-try:
-    pipe_txt2img = StableDiffusionPipeline.from_pretrained(
-        MODEL_ID,
-        revision=REVISION,
-        torch_dtype=dtype,
-        safety_checker=None,
-    ).to(device)
-
-    # 🔑 FORCE tokenizer + text encoder
-    pipe_txt2img.tokenizer = CLIPTokenizer.from_pretrained(
-        MODEL_ID, subfolder="tokenizer"
-    )
-    pipe_txt2img.text_encoder = CLIPTextModel.from_pretrained(
-        MODEL_ID,
-        subfolder="text_encoder",
-        torch_dtype=dtype,
-    ).to(device)
-
-    # Scheduler
-    pipe_txt2img.scheduler = EulerAncestralDiscreteScheduler.from_config(
-        pipe_txt2img.scheduler.config
-    )
-
-    # Memory optimisations
-    try:
-        pipe_txt2img.enable_attention_slicing()
-        pipe_txt2img.enable_vae_slicing()
-    except Exception:
-        pass
-
-    try:
-        pipe_txt2img.enable_xformers_memory_efficient_attention()
-    except Exception:
-        pass
-
-    pipe_txt2img.set_progress_bar_config(disable=True)
-
-    # Img2Img pipeline (share components)
-    pipe_img2img = StableDiffusionImg2ImgPipeline(**pipe_txt2img.components).to(device)
-    pipe_img2img.scheduler = EulerAncestralDiscreteScheduler.from_config(
-        pipe_img2img.scheduler.config
-    )
+pipe_txt2img = StableDiffusionPipeline.from_pretrained(
+    MODEL_ID,
+    revision=REVISION,
+    torch_dtype=dtype,
+    safety_checker=None,
+).to(device)
+
+# 🔑 Force tokenizer + text encoder (fixes tokenize None bug)
+pipe_txt2img.tokenizer = CLIPTokenizer.from_pretrained(
+    MODEL_ID, subfolder="tokenizer"
+)
+pipe_txt2img.text_encoder = CLIPTextModel.from_pretrained(
+    MODEL_ID,
+    subfolder="text_encoder",
+    torch_dtype=dtype,
+).to(device)
+
+# Scheduler
+pipe_txt2img.scheduler = EulerAncestralDiscreteScheduler.from_config(
+    pipe_txt2img.scheduler.config
+)
 
-    # Defensive checks
-    assert pipe_txt2img.tokenizer is not None
-    assert pipe_txt2img.text_encoder is not None
+# Memory optimisations (safe on Spaces)
+pipe_txt2img.enable_attention_slicing()
+pipe_txt2img.enable_vae_slicing()
 
-    model_loaded = True
+try:
+    pipe_txt2img.enable_xformers_memory_efficient_attention()
+except Exception:
+    pass
 
-except Exception as e:
-    load_error = repr(e)
-    model_loaded = False
+pipe_txt2img.set_progress_bar_config(disable=True)
 
-# ============================================================
-# Helpers
-# ============================================================
-def _make_error_image(w, h):
-    return Image.new("RGB", (w, h), (30, 30, 40))
+# Img2Img pipeline (reuse components)
+pipe_img2img = StableDiffusionImg2ImgPipeline(**pipe_txt2img.components).to(device)
+pipe_img2img.scheduler = EulerAncestralDiscreteScheduler.from_config(
+    pipe_img2img.scheduler.config
+)
 
 # ============================================================
-# Inference
+# GPU INFERENCE FUNCTION (Spaces requires this)
 # ============================================================
+@spaces.GPU
 def infer(
     prompt,
     negative_prompt,
@@ -124,9 +104,6 @@ def infer(
     width = int(width)
     height = int(height)
 
-    if not model_loaded:
-        return _make_error_image(width, height), load_error
-
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
@@ -157,22 +134,15 @@ def infer(
 
         return image, f"Seed: {seed}"
 
-    except Exception as e:
-        return _make_error_image(width, height), str(e)
-
     finally:
         gc.collect()
-        if device.type == "cuda":
-            torch.cuda.empty_cache()
+        torch.cuda.empty_cache()
 
 # ============================================================
 # UI
 # ============================================================
 with gr.Blocks(title="Stable Diffusion (Unlearning Model)") as demo:
-    gr.Markdown("## Stable Diffusion Generator")
-
-    if not model_loaded:
-        gr.Markdown(f"⚠️ **Model failed to load**\n\n{load_error}")
+    gr.Markdown("## Stable Diffusion Generator (GPU)")
 
     prompt = gr.Textbox(label="Prompt", lines=2)
     init_image = gr.Image(label="Initial image (optional)", type="pil")
@@ -208,4 +178,4 @@ with gr.Blocks(title="Stable Diffusion (Unlearning Model)") as demo:
         outputs=[result, status],
     )
 
-demo.queue().launch(ssr_mode=False)
+demo.queue().launch(server_name="0.0.0.0", server_port=7860)