Update app.py

app.py

@@ -10,22 +10,11 @@ import torch
 from diffusers import (
     StableDiffusionPipeline,
     StableDiffusionImg2ImgPipeline,
-    StableDiffusionXLPipeline,
-    StableDiffusionXLImg2ImgPipeline,
     EulerAncestralDiscreteScheduler,
 )
+from transformers import CLIPTokenizer, CLIPTextModel
 from huggingface_hub import login
 
-# ============================================================
-# Optional GPU decorator (Spaces)
-# ============================================================
-try:
-    import spaces
-    GPU_DECORATOR = spaces.GPU
-except Exception:
-    def GPU_DECORATOR(fn):
-        return fn
-
 # ============================================================
 # Config
 # ============================================================
@@ -41,42 +30,38 @@ device = torch.device("cuda" if cuda_available else "cpu")
 dtype = torch.float16 if cuda_available else torch.float32
 
 MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1216 if cuda_available else 768
+MAX_IMAGE_SIZE = 768 if not cuda_available else 1024
 
 pipe_txt2img = None
 pipe_img2img = None
-is_sdxl = False
 model_loaded = False
 load_error = None
 
 # ============================================================
-# Load model (AUTO detect SDXL vs SD)
+# Load model (FORCED tokenizer fix)
 # ============================================================
 try:
-    from_pretrained_kwargs = dict(
-        torch_dtype=dtype,
+    pipe_txt2img = StableDiffusionPipeline.from_pretrained(
+        MODEL_ID,
         revision=REVISION,
-    )
-
-    if HF_TOKEN:
-        from_pretrained_kwargs["token"] = HF_TOKEN
+        torch_dtype=dtype,
+        safety_checker=None,
+    ).to(device)
 
-    # Try SDXL first
-    try:
-        pipe_txt2img = StableDiffusionXLPipeline.from_pretrained(
-            MODEL_ID, **from_pretrained_kwargs
-        )
-        is_sdxl = True
-    except Exception:
-        pipe_txt2img = StableDiffusionPipeline.from_pretrained(
-            MODEL_ID, **from_pretrained_kwargs
-        )
-        is_sdxl = False
+    # 🔑 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
     )
-    pipe_txt2img = pipe_txt2img.to(device)
 
     # Memory optimisations
     try:
@@ -92,16 +77,15 @@ try:
 
     pipe_txt2img.set_progress_bar_config(disable=True)
 
-    # Create img2img pipeline
-    if is_sdxl:
-        pipe_img2img = StableDiffusionXLImg2ImgPipeline(**pipe_txt2img.components)
-    else:
-        pipe_img2img = StableDiffusionImg2ImgPipeline(**pipe_txt2img.components)
-
+    # Img2Img pipeline (share components)
+    pipe_img2img = StableDiffusionImg2ImgPipeline(**pipe_txt2img.components).to(device)
     pipe_img2img.scheduler = EulerAncestralDiscreteScheduler.from_config(
         pipe_img2img.scheduler.config
     )
-    pipe_img2img = pipe_img2img.to(device)
+
+    # Defensive checks
+    assert pipe_txt2img.tokenizer is not None
+    assert pipe_txt2img.text_encoder is not None
 
     model_loaded = True
 
@@ -112,13 +96,12 @@ except Exception as e:
 # ============================================================
 # Helpers
 # ============================================================
-def _make_error_image(w, h, text):
+def _make_error_image(w, h):
     return Image.new("RGB", (w, h), (30, 30, 40))
 
 # ============================================================
 # Inference
 # ============================================================
-@GPU_DECORATOR
 def infer(
     prompt,
     negative_prompt,
@@ -135,19 +118,13 @@ def infer(
     height = int(height)
 
     if not model_loaded:
-        return _make_error_image(width, height, "Model not loaded"), load_error
+        return _make_error_image(width, height), load_error
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator(device=device).manual_seed(seed)
 
-    common_kwargs = dict(
-        guidance_scale=float(guidance_scale),
-        num_inference_steps=int(num_inference_steps),
-        generator=generator,
-    )
-
     try:
         with torch.inference_mode():
             if init_image is not None:
@@ -156,7 +133,9 @@ def infer(
                     negative_prompt=negative_prompt,
                     image=init_image,
                     strength=float(strength),
-                    **common_kwargs,
+                    guidance_scale=float(guidance_scale),
+                    num_inference_steps=int(num_inference_steps),
+                    generator=generator,
                 ).images[0]
             else:
                 image = pipe_txt2img(
@@ -164,13 +143,15 @@ def infer(
                     negative_prompt=negative_prompt,
                     width=width,
                     height=height,
-                    **common_kwargs,
+                    guidance_scale=float(guidance_scale),
+                    num_inference_steps=int(num_inference_steps),
+                    generator=generator,
                 ).images[0]
 
-        return image, f"Seed: {seed} | {'SDXL' if is_sdxl else 'SD 1.x'}"
+        return image, f"Seed: {seed}"
 
     except Exception as e:
-        return _make_error_image(width, height, "Generation error"), str(e)
+        return _make_error_image(width, height), str(e)
 
     finally:
         gc.collect()
@@ -180,8 +161,7 @@ def infer(
 # ============================================================
 # UI
 # ============================================================
-with gr.Blocks(title="Text-to-Image / Image-to-Image") as demo:
-
+with gr.Blocks(title="Stable Diffusion (Unlearning Model)") as demo:
     gr.Markdown("## Stable Diffusion Generator")
 
     if not model_loaded:
@@ -196,11 +176,11 @@ with gr.Blocks(title="Text-to-Image / Image-to-Image") as demo:
 
     with gr.Accordion("Advanced Settings", open=False):
         negative_prompt = gr.Textbox(label="Negative prompt", value="")
-        seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
+        seed = gr.Slider(0, MAX_SEED, step=1, value=0, label="Seed")
         randomize_seed = gr.Checkbox(True, label="Randomize seed")
         width = gr.Slider(256, MAX_IMAGE_SIZE, step=32, value=512, label="Width")
         height = gr.Slider(256, MAX_IMAGE_SIZE, step=32, value=512, label="Height")
-        guidance_scale = gr.Slider(0, 20, step=0.1, value=7.5, label="Guidance scale")
+        guidance_scale = gr.Slider(1, 20, step=0.5, value=7.5, label="Guidance scale")
         num_inference_steps = gr.Slider(1, 40, step=1, value=20, label="Steps")
         strength = gr.Slider(0.0, 1.0, step=0.05, value=0.7, label="Image strength")