Update app.py

app.py

@@ -5,7 +5,7 @@
 # - @spaces.GPU MUST be used directly
 # ============================================================
 
-import spaces  # MUST be first, no try/except
+import spaces  # MUST be first
 
 import os
 import random
@@ -37,11 +37,8 @@ if HF_TOKEN:
 device = torch.device("cuda")
 dtype = torch.float16
 
+IMAGE_SIZE = 512
 MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-pipe_txt2img = None
-pipe_img2img = None
 
 # ============================================================
 # Load model (once at startup)
@@ -53,7 +50,7 @@ pipe_txt2img = StableDiffusionPipeline.from_pretrained(
     safety_checker=None,
 ).to(device)
 
-# 🔑 Force tokenizer + text encoder (fixes tokenize None bug)
+# 🔑 Force tokenizer + text encoder
 pipe_txt2img.tokenizer = CLIPTokenizer.from_pretrained(
     MODEL_ID, subfolder="tokenizer"
 )
@@ -68,7 +65,7 @@ pipe_txt2img.scheduler = EulerAncestralDiscreteScheduler.from_config(
     pipe_txt2img.scheduler.config
 )
 
-# Memory optimisations (safe on Spaces)
+# Memory optimisations
 pipe_txt2img.enable_attention_slicing()
 pipe_txt2img.enable_vae_slicing()
 
@@ -79,8 +76,10 @@ except Exception:
 
 pipe_txt2img.set_progress_bar_config(disable=True)
 
-# Img2Img pipeline (reuse components)
-pipe_img2img = StableDiffusionImg2ImgPipeline(**pipe_txt2img.components).to(device)
+# Img2Img pipeline
+pipe_img2img = StableDiffusionImg2ImgPipeline(
+    **pipe_txt2img.components
+).to(device)
 pipe_img2img.scheduler = EulerAncestralDiscreteScheduler.from_config(
     pipe_img2img.scheduler.config
 )
@@ -94,16 +93,11 @@ def infer(
     negative_prompt,
     seed,
     randomize_seed,
-    width,
-    height,
     guidance_scale,
     num_inference_steps,
     init_image,
     strength,
 ):
-    width = int(width)
-    height = int(height)
-
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
@@ -125,8 +119,8 @@ def infer(
                 image = pipe_txt2img(
                     prompt=prompt,
                     negative_prompt=negative_prompt,
-                    width=width,
-                    height=height,
+                    width=IMAGE_SIZE,
+                    height=IMAGE_SIZE,
                     guidance_scale=float(guidance_scale),
                     num_inference_steps=int(num_inference_steps),
                     generator=generator,
@@ -141,25 +135,34 @@ def infer(
 # ============================================================
 # UI
 # ============================================================
-with gr.Blocks(title="Stable Diffusion (Unlearning Model)") as demo:
-    gr.Markdown("## Stable Diffusion Generator (GPU)")
+with gr.Blocks(title="Stable Diffusion (512×512)") as demo:
+    gr.Markdown("## Stable Diffusion Generator (GPU, 512×512)")
 
-    prompt = gr.Textbox(label="Prompt", lines=2)
-    init_image = gr.Image(label="Initial image (optional)", type="pil")
+    prompt = gr.Textbox(
+        label="Prompt",
+        placeholder="Describe the image you want to generate",
+        lines=2,
+    )
+
+    init_image = gr.Image(
+        label="Initial image (optional, enables img2img)",
+        type="pil",
+    )
 
     run_button = gr.Button("Generate")
     result = gr.Image(label="Result")
     status = gr.Markdown("")
 
     with gr.Accordion("Advanced Settings", open=False):
-        negative_prompt = gr.Textbox(label="Negative prompt", value="nsfw, (low quality, worst quality:1.2), very displeasing, 3d, watermark, signature, ugly, poorly drawn, deform, paining, eyes misaligned, bad,  disproportional, boring,  fuzzy, asymmetric")
+        negative_prompt = gr.Textbox(
+            label="Negative prompt",
+            value="nsfw, (low quality, worst quality:1.2), watermark, signature, ugly, deformed",
+        )
         seed = gr.Slider(0, MAX_SEED, step=1, value=0, label="Seed")
         randomize_seed = gr.Checkbox(True, label="Randomize seed")
-        width = 512 #gr.Slider(256, MAX_IMAGE_SIZE, step=32, value=512, label="Width")
-        height =  512 # gr.Slider(256, MAX_IMAGE_SIZE, step=32, value=512, label="Height")
         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=30, label="Steps")
-        strength = gr.Slider(0.0, 1.0, step=0.05, value=0.7, label="Image strength")
+        strength = gr.Slider(0.0, 1.0, step=0.05, value=0.7, label="Image strength (img2img)")
 
     run_button.click(
         fn=infer,
@@ -168,8 +171,6 @@ with gr.Blocks(title="Stable Diffusion (Unlearning Model)") as demo:
             negative_prompt,
             seed,
             randomize_seed,
-            width,
-            height,
             guidance_scale,
             num_inference_steps,
             init_image,