Update app.py

app.py

@@ -2,62 +2,73 @@ import torch
 import gradio as gr
 from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline
 from PIL import Image
-import numpy as np
 from huggingface_hub import login
 
-# Authenticate (asks only once)
-hf_token = input("🔐 Enter your Hugging Face token: ").strip()
-login(hf_token)
-
-# Check device
+# Globals
+text2img_pipe = None
+inpaint_pipe = None
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load models
-text2img_pipe = StableDiffusionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-2-1",
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-    use_auth_token=hf_token
-).to(device)
+# Text-to-image generation
+def generate_image(hf_token, prompt):
+    global text2img_pipe
+    try:
+        if text2img_pipe is None:
+            login(hf_token)
+            text2img_pipe = StableDiffusionPipeline.from_pretrained(
+                "stabilityai/stable-diffusion-2-1",
+                torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+                use_auth_token=hf_token
+            ).to(device)
 
-inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
-    "stable-diffusion-v1-5/stable-diffusion-inpainting",
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-    use_auth_token=hf_token
-).to(device)
+        image = text2img_pipe(prompt).images[0]
+        return "✅ Image generated", image
+    except Exception as e:
+        return f"❌ Error: {str(e)}", None
 
-# Image generation
-def generate_image(prompt):
-    image = text2img_pipe(prompt).images[0]
-    return image
+# Image manipulation (inpainting)
+def manipulate_image(hf_token, prompt, image, mask):
+    global inpaint_pipe
+    try:
+        if image is None or mask is None:
+            return "⚠️ Please upload both image and mask", None
 
-# Image inpainting (manipulation)
-def manipulate_image(prompt, image, mask):
-    if image is None or mask is None:
-        return "Please upload both image and mask.", None
+        if inpaint_pipe is None:
+            login(hf_token)
+            inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
+                "stable-diffusion-v1-5/stable-diffusion-inpainting",
+                torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+                use_auth_token=hf_token
+            ).to(device)
 
-    image = image.convert("RGB").resize((512, 512))
-    mask = mask.convert("L").resize((512, 512))  # grayscale
+        image = image.convert("RGB").resize((512, 512))
+        mask = mask.convert("L").resize((512, 512))
 
-    result = inpaint_pipe(prompt=prompt, image=image, mask_image=mask).images[0]
-    return "Image updated!", result
+        result = inpaint_pipe(prompt=prompt, image=image, mask_image=mask).images[0]
+        return "✅ Image updated", result
+    except Exception as e:
+        return f"❌ Error: {str(e)}", None
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🎨 AI Image Generator & Editor")
+    gr.Markdown("## 🎨 AI Image Generator & Editor (Stable Diffusion)")
 
     with gr.Tab("🖼️ Generate Image from Text"):
-        prompt = gr.Textbox(label="Describe your image")
-        gen_btn = gr.Button("Generate")
+        token1 = gr.Textbox(label="🔐 Hugging Face Token", type="password")
+        prompt1 = gr.Textbox(label="Prompt (e.g., 'a sunset over mountains')")
+        status1 = gr.Textbox(label="Status")
         gen_output = gr.Image(label="Generated Image")
-        gen_btn.click(fn=generate_image, inputs=prompt, outputs=gen_output)
+        gen_btn = gr.Button("Generate Image")
+        gen_btn.click(fn=generate_image, inputs=[token1, prompt1], outputs=[status1, gen_output])
 
-    with gr.Tab("🪄 Manipulate Existing Image"):
-        prompt2 = gr.Textbox(label="Describe the edit (e.g., 'remove the tree')")
+    with gr.Tab("🪄 Manipulate Uploaded Image"):
+        token2 = gr.Textbox(label="🔐 Hugging Face Token", type="password")
+        prompt2 = gr.Textbox(label="Edit Prompt (e.g., 'remove the background')")
         img_input = gr.Image(label="Upload Image", type="pil", tool="editor")
-        mask_input = gr.Image(label="Draw Mask on Area to Edit", tool="sketch", type="pil")
-        edit_btn = gr.Button("Apply Edit")
-        status = gr.Textbox(label="Status")
+        mask_input = gr.Image(label="Draw Mask", type="pil", tool="sketch")
+        status2 = gr.Textbox(label="Status")
         edited_img = gr.Image(label="Edited Image")
-        edit_btn.click(fn=manipulate_image, inputs=[prompt2, img_input, mask_input], outputs=[status, edited_img])
+        edit_btn = gr.Button("Apply Edit")
+        edit_btn.click(fn=manipulate_image, inputs=[token2, prompt2, img_input, mask_input], outputs=[status2, edited_img])
 
-demo.launch(server_name="0.0.0.0", server_port=7860)
+demo.launch()