vb

app.py

@@ -1,24 +1,10 @@
 import os
 import torch
 import gradio as gr
-from diffusers import DiffusionPipeline
+from diffusers import StableDiffusionImg2ImgPipeline
 from PIL import Image
 
-# --- Configuration ---
-SPACE_TITLE = "🎨 Enhanced Studio Ghibli AI Art Generator (LoRA)"
-SPACE_DESCRIPTION = "Upload a portrait or a photo and transform it into a breathtaking Studio Ghibli-style masterpiece using a LoRA for fine-tuned results."
-BASE_MODEL_ID = "black-forest-labs/FLUX.1-dev"
-LORA_REPO_ID = "strangerzonehf/Flux-Ghibli-Art-LoRA"
-TRIGGER_WORD = "Ghibli Art"
-STRENGTH = 0.60  # Adjust for better balance between input and style
-GUIDANCE_SCALE = 7.5  # Increased for better prompt adherence
-NUM_INFERENCE_STEPS = 30  # Increased for potentially higher quality
-INPUT_IMAGE_SIZE = (512, 512)
-PROMPT_PREFIX = ""  # No need for separate prefix as LoRA is targeted
-NEGATIVE_PROMPT = "ugly, deformed, blurry, low quality, bad anatomy, bad proportions, disfigured, poorly drawn face, mutation, mutated, extra limbs, extra fingers, body horror, glitchy, tiling"
-
-# --- Device Setup ---
-# Force CUDA usage, assuming A100 is the first GPU (index 0)
+# Force CUDA usage if available
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 torch.backends.cudnn.benchmark = True
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -32,87 +18,61 @@ try:
     torch.zeros(1).to(device)
     print("✅ Torch initialized successfully on", device)
 except Exception as e:
-    print(f"⚠️ Torch initialization error: {e}")
+    print("⚠️ Torch initialization error:", e)
 
-# --- Model Loading ---
-try:
-    pipe = DiffusionPipeline.from_pretrained(BASE_MODEL_ID, torch_dtype=torch.bfloat16)
-except ValueError as e:
-    if "sentencepiece" in str(e):
-        print("⚠️ Error: sentencepiece is not installed. Please install it with: pip install sentencepiece")
-        raise
-    else:
-        raise e
+# Load the correct Stable Diffusion pipeline
+model_id = "nitrosocke/Ghibli-Diffusion"
+
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+    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.load_lora_weights(LORA_REPO_ID)
-    print(f"✅ LoRA weights loaded from {LORA_REPO_ID}")
+    pipe.enable_xformers_memory_efficient_attention()
+    print("✅ xFormers enabled!")
 except Exception as e:
-    print(f"⚠️ Error loading LoRA weights: {e}")
+    print(f"⚠️ xFormers not available: {e}")
 
-pipe.to(device)
-
-# --- 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()
+# Apply additional optimizations for performance
+pipe.enable_model_cpu_offload()
 pipe.enable_vae_slicing()
 pipe.enable_attention_slicing()
 
-# --- Image Transformation Function ---
-def transform_image(input_image):
-    if input_image is None:
-        return None
-
-    try:
-        input_image = input_image.resize(INPUT_IMAGE_SIZE)
+# Enhanced Studio Ghibli-style transformation prompt
+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."
+)
 
-        prompt = f"{PROMPT_PREFIX} {TRIGGER_WORD}, portrait of a person" # Incorporate trigger word
+# Image transformation function
+def transform_image(input_image):
+    input_image = input_image.resize((512, 512))
 
-        output = pipe(
-            prompt=prompt,
-            image=input_image,
-            strength=STRENGTH,
-            guidance_scale=GUIDANCE_SCALE,
-            num_inference_steps=NUM_INFERENCE_STEPS,
-            negative_prompt=NEGATIVE_PROMPT,
-        ).images[0]
+    # Pass the image as `init_image`
+    output = pipe(
+        prompt=prompt,
+        init_image=input_image,  # ✅ FIXED: Changed from "image" to "init_image"
+        strength=0.65,
+        guidance_scale=5.0,  # Slightly increased for better stylization
+        num_inference_steps=25,  # More steps for higher quality output
+    )
 
-        return output
-    except Exception as e:
-        print(f"❌ Error during image transformation: {e}")
-        return None
+    return output.images[0]
 
-# --- Gradio UI ---
-iface = gr.Interface(
+# Gradio UI
+demo = 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=SPACE_TITLE,
-    description=SPACE_DESCRIPTION,
-    examples=[
-        "examples/portrait1.jpg",
-        "examples/photo1.jpg",
-        "examples/landscape1.jpg",
-    ],
+    title="🎨 Studio Ghibli AI Art Generator",
+    description="Upload a portrait or a photo and transform it into a breathtaking Studio Ghibli-style masterpiece!",
 )
 
-# --- Main Execution ---
 if __name__ == "__main__":
-    # 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()
+    demo.launch()