dcsdcv

app.py

@@ -4,7 +4,17 @@ import gradio as gr
 from diffusers import StableDiffusionImg2ImgPipeline
 from PIL import Image
 
-# Force CUDA usage if available
+# --- Configuration ---
+SPACE_TITLE = "🎨 Studio Ghibli AI Art Generator"
+SPACE_DESCRIPTION = "Upload a portrait or a photo and transform it into a breathtaking Studio Ghibli-style masterpiece!"
+MODEL_ID = "nitrosocke/Ghibli-Diffusion"
+STRENGTH = 0.65
+GUIDANCE_SCALE = 5.0
+NUM_INFERENCE_STEPS = 25
+INPUT_IMAGE_SIZE = (512, 512)
+
+# --- Device Setup ---
+# Force CUDA usage, assuming A100 is the first GPU (index 0)
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 torch.backends.cudnn.benchmark = True
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -18,60 +28,82 @@ try:
     torch.zeros(1).to(device)
     print("✅ Torch initialized successfully on", device)
 except Exception as e:
-    print("⚠️ Torch initialization error:", e)
-
-# Load the optimized Stable Diffusion model
-model_id = "nitrosocke/Ghibli-Diffusion"
+    print(f"⚠️ Torch initialization error: {e}")
 
+# --- Model Loading ---
 pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-    model_id,
+    MODEL_ID,
     torch_dtype=torch.float16 if device == "cuda" else torch.float32,
     use_safetensors=True,
     low_cpu_mem_usage=True
 ).to(device)
 
-# Try enabling xFormers for memory efficiency
-try:
-    pipe.enable_xformers_memory_efficient_attention()
-    print("✅ xFormers enabled!")
-except Exception as e:
-    print(f"⚠️ xFormers not available: {e}")
-
-# Apply additional optimizations for performance
-pipe.enable_model_cpu_offload()
+# --- Optimization (Conditional for CUDA) ---
+if device == "cuda":
+    try:
+        pipe.enable_xformers_memory_efficient_attention()
+        print("✅ xFormers enabled!")
+    except Exception as e:
+        print(f"⚠️ xFormers not available: {e}")
+    pipe.enable_model_cpu_offload() # Keep most of the model on GPU, offload selectively
 pipe.enable_vae_slicing()
 pipe.enable_attention_slicing()
 
-# Enhanced Studio Ghibli-style transformation prompt
-prompt = (
+# --- Prompt Definition ---
+GHIBLI_PROMPT = (
     "Studio Ghibli anime-style illustration, magical landscape, soft pastel colors, "
     "hand-painted textures, cinematic lighting, dreamy atmosphere, vibrant and rich details, "
     "Miyazaki-inspired fantasy world, watercolor aesthetic, warm sunlight, intricate composition, "
     "high detail, whimsical and nostalgic beauty."
 )
 
-# Image transformation function
+# --- Image Transformation Function ---
 def transform_image(input_image):
-    input_image = input_image.resize((512, 512))
+    if input_image is None:
+        return None
+
+    try:
+        input_image = input_image.resize(INPUT_IMAGE_SIZE)
 
-    output = pipe(
-        prompt=prompt,
-        image=input_image,
-        strength=0.65,
-        guidance_scale=5.0,  # Slightly increased for better stylization
-        num_inference_steps=25,  # More steps for higher quality output
-    )
+        output = pipe(
+            prompt=GHIBLI_PROMPT,
+            image=input_image,
+            strength=STRENGTH,
+            guidance_scale=GUIDANCE_SCALE,
+            num_inference_steps=NUM_INFERENCE_STEPS,
+        )
 
-    return output.images[0]
+        return output.images[0]
+    except Exception as e:
+        print(f"❌ Error during image transformation: {e}")
+        return None
 
-# Gradio UI
-demo = gr.Interface(
+# --- Gradio UI ---
+iface = gr.Interface(
     fn=transform_image,
     inputs=gr.Image(type="pil", label="Upload a Portrait/Photo"),
     outputs=gr.Image(type="pil", label="Studio Ghibli-Style Output"),
-    title="🎨 Studio Ghibli AI Art Generator",
-    description="Upload a portrait or a photo and transform it into a breathtaking Studio Ghibli-style masterpiece!",
+    title=SPACE_TITLE,
+    description=SPACE_DESCRIPTION,
+    examples=[
+        "examples/portrait1.jpg",
+        "examples/photo1.jpg",
+        "examples/landscape1.jpg",
+    ],
 )
 
+# --- Main Execution ---
 if __name__ == "__main__":
-    demo.launch()
+    # Create an 'examples' directory if it doesn't exist and add some sample images
+    if not os.path.exists("examples"):
+        os.makedirs("examples")
+        # You'll need to download or create these example images
+        # and place them in the 'examples' folder.
+        # Example:
+        # from urllib.request import urlretrieve
+        # urlretrieve("URL_TO_YOUR_EXAMPLE_IMAGE_1", "examples/portrait1.jpg")
+        # urlretrieve("URL_TO_YOUR_EXAMPLE_IMAGE_2", "examples/photo1.jpg")
+        # urlretrieve("URL_TO_YOUR_EXAMPLE_IMAGE_3", "examples/landscape1.jpg")
+        print("ℹ️ Created 'examples' directory. Please add sample images.")
+
+    iface.launch()