Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
 from PIL import Image
 import io
 import requests
@@ -11,8 +11,8 @@ import tempfile
 import time
 import base64
 
-# External OCI API URL - THIS IS THE KEY FIX
-OCI_API_BASE_URL = "https://yukee1992-oci-story-book.hf.space"
+# External OCI API URL
+OCI_API_BASE_URL = "https://yukee1992-oci-video-storage.hf.space"
 
 # Try to import your existing OCI connector for direct access
 try:
@@ -23,94 +23,103 @@ except ImportError:
     DIRECT_OCI_ACCESS = False
     print("⚠️  Direct OCI access not available - using external API endpoint")
 
-# Initialize the Stable Diffusion model
-def load_model():
-    """Load and return the Stable Diffusion model"""
-    print("🔄 Loading Stable Diffusion model...")
+# HIGH-QUALITY MODEL SELECTION
+MODEL_CHOICES = {
+    "dreamshaper-8": "lykon/dreamshaper-8",  # Excellent for fantasy/artistic
+    "realistic-vision": "SG161222/Realistic_Vision_V5.1",  # Photorealistic
+    "anything-v5": "andite/anything-v5.0",  # Anime/illustration style
+    "openjourney": "prompthero/openjourney",  # Artistic/Midjourney style
+    "sd-2.1": "stabilityai/stable-diffusion-2-1",  # General purpose
+}
+
+# Initialize the HIGH-QUALITY Stable Diffusion model
+def load_model(model_name="dreamshaper-8"):
+    """Load and return a high-quality Stable Diffusion model"""
+    print(f"🔄 Loading {model_name} model...")
     try:
+        model_id = MODEL_CHOICES.get(model_name, "lykon/dreamshaper-8")
+        
         pipe = StableDiffusionPipeline.from_pretrained(
-            "runwayml/stable-diffusion-v1-5", 
+            model_id, 
             torch_dtype=torch.float32,
             safety_checker=None,
             requires_safety_checker=False
-        ).to("cpu")
-        print("✅ Model loaded successfully!")
+        )
+        
+        # Use better scheduler for quality
+        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+        pipe = pipe.to("cpu")
+        
+        print(f"✅ {model_name} loaded successfully!")
         return pipe
     except Exception as e:
         print(f"❌ Model loading failed: {e}")
-        raise e
+        # Fallback to SD 1.5
+        return StableDiffusionPipeline.from_pretrained(
+            "runwayml/stable-diffusion-v1-5", 
+            torch_dtype=torch.float32,
+            safety_checker=None,
+            requires_safety_checker=False
+        ).to("cpu")
 
 # Load the model once at startup
 pipe = load_model()
 
-def verify_oci_upload(object_name):
-    """Verify that the file actually exists in OCI"""
-    try:
-        if DIRECT_OCI_ACCESS:
-            # Use direct OCI client to check if file exists
-            namespace = oci_connector.config.get('namespace')
-            bucket_name = oci_connector.config.get('bucket_name')
-            
-            if not namespace or not bucket_name:
-                return False, "Namespace or bucket name not configured"
-            
-            try:
-                response = oci_connector.client.head_object(
-                    namespace_name=namespace,
-                    bucket_name=bucket_name,
-                    object_name=object_name
-                )
-                return True, "✅ File verified in OCI"
-            except oci.exceptions.ServiceError as e:
-                if e.status == 404:
-                    return False, "❌ File not found in OCI (404)"
-                else:
-                    return False, f"❌ OCI Error: {e.message}"
-        
-        else:
-            # For API mode, we can't easily verify, so just return optimistic
-            return True, "✅ Upload reported successful (verification not available)"
-            
-    except Exception as e:
-        return False, f"❌ Verification failed: {str(e)}"
-
-def save_to_oci_direct(image, prompt):
-    """Save image using direct OCI connector access with verification"""
-    try:
-        # Create temporary file
-        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
-            image.save(tmp, format='PNG')
-            temp_path = tmp.name
-        
-        # Create organized filename
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        safe_prompt = "".join(c for c in prompt[:30] if c.isalnum() or c in (' ', '-', '_')).rstrip()
-        filename = f"story_{timestamp}_{safe_prompt}.png"
-        
-        # Use project-based directory structure
-        object_name = f"storybook-generator/childrens-books/{filename}"
-        
-        print(f"📤 Uploading to OCI: {object_name}")
-        
-        # Upload using existing OCI connector
-        success, message = oci_connector.upload_file(temp_path, object_name, None)
-        
-        # Clean up temporary file
-        os.unlink(temp_path)
-        
-        if success:
-            # Verify the upload actually worked
-            print("🔍 Verifying OCI upload...")
-            verified, verify_msg = verify_oci_upload(object_name)
-            if verified:
-                return f"✅ {message} | {verify_msg}"
-            else:
-                return f"⚠️  {message} BUT {verify_msg}"
-        else:
-            return f"❌ {message}"
-            
-    except Exception as e:
-        return f"❌ Direct upload failed: {str(e)}"
+# PROFESSIONAL PROMPT ENGINEERING
+def enhance_prompt(prompt, style="childrens_book"):
+    """Transform basic prompts into professional-grade prompts"""
+    
+    style_templates = {
+        "childrens_book": [
+            "masterpiece, best quality, 4K, ultra detailed, children's book illustration",
+            "watercolor painting, whimsical, cute, charming, storybook style",
+            "vibrant colors, soft lighting, magical, enchanting, dreamlike",
+            "Pixar style, Disney animation, high detail, professional artwork"
+        ],
+        "realistic": [
+            "photorealistic, 8K, ultra detailed, professional photography",
+            "sharp focus, studio lighting, high resolution, intricate details",
+            "realistic textures, natural lighting, cinematic quality"
+        ],
+        "fantasy": [
+            "epic fantasy art, digital painting, concept art, trending on artstation",
+            "magical, mystical, ethereal, otherworldly, fantasy illustration",
+            "dynamic composition, dramatic lighting, highly detailed"
+        ],
+        "anime": [
+            "anime style, Japanese animation, high quality, detailed artwork",
+            "beautiful anime illustration, vibrant colors, clean lines",
+            "studio ghibli style, makoto shinkai, professional anime art"
+        ]
+    }
+    
+    # Choose style template
+    templates = style_templates.get(style, style_templates["childrens_book"])
+    style_prompt = templates[0]  # Use the first template
+    
+    # Enhanced prompt construction
+    enhanced = f"{style_prompt}, {prompt}"
+    
+    # Add quality boosters
+    quality_boosters = [
+        "intricate details", "beautiful composition", "perfect lighting",
+        "professional artwork", "award winning", "trending on artstation"
+    ]
+    
+    # Add 2-3 random quality boosters
+    import random
+    boosters = random.sample(quality_boosters, 2)
+    enhanced += ", " + ", ".join(boosters)
+    
+    # Negative prompt to avoid bad quality
+    negative_prompt = (
+        "blurry, low quality, low resolution, ugly, deformed, poorly drawn, "
+        "bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, "
+        "disconnected limbs, mutation, mutated, ugly, disgusting, bad art, "
+        "beginner, amateur, distorted, watermark, signature, text, username"
+    )
+    
+    return enhanced, negative_prompt
 
 def save_to_oci_via_api(image, prompt):
     """Save image using the EXTERNAL OCI API endpoint"""
@@ -123,9 +132,9 @@ def save_to_oci_via_api(image, prompt):
         # Create filename
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         safe_prompt = "".join(c for c in prompt[:30] if c.isalnum() or c in (' ', '-', '_')).rstrip()
-        filename = f"story_{timestamp}_{safe_prompt}.png"
+        filename = f"hq_story_{timestamp}_{safe_prompt}.png"
         
-        # Use the EXTERNAL API URL - THIS IS THE FIX
+        # Use the EXTERNAL API URL
         api_url = f"{OCI_API_BASE_URL}/api/upload"
         
         print(f"🌐 Calling EXTERNAL OCI API: {api_url}")
@@ -137,15 +146,13 @@ def save_to_oci_via_api(image, prompt):
         
         data = {
             'project_id': 'storybook-generator',
-            'subfolder': 'childrens-books'
+            'subfolder': 'high-quality-books'
         }
         
         # Make the API request with timeout
         response = requests.post(api_url, files=files, data=data, timeout=30)
         
         print(f"📨 API Response: {response.status_code}")
-        if response.text:
-            print(f"📨 API Response Text: {response.text[:200]}...")
         
         if response.status_code == 200:
             result = response.json()
@@ -156,188 +163,133 @@ def save_to_oci_via_api(image, prompt):
         else:
             return f"❌ HTTP Error: {response.status_code} - {response.text}"
             
-    except requests.exceptions.Timeout:
-        return "❌ API call timed out (30s)"
-    except requests.exceptions.ConnectionError:
-        return "❌ Cannot connect to OCI API"
     except Exception as e:
         return f"❌ API upload failed: {str(e)}"
 
-def save_to_oci_via_api_text(image, prompt):
-    """Alternative: Save image using the external OCI API text endpoint"""
-    try:
-        # Convert image to base64
-        img_bytes = io.BytesIO()
-        image.save(img_bytes, format='PNG')
-        img_base64 = base64.b64encode(img_bytes.getvalue()).decode('utf-8')
-        
-        # Create filename
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        safe_prompt = "".join(c for c in prompt[:30] if c.isalnum() or c in (' ', '-', '_')).rstrip()
-        filename = f"story_{timestamp}_{safe_prompt}.png"
-        
-        # Use the EXTERNAL API text endpoint
-        api_url = f"{OCI_API_BASE_URL}/api/upload-text"
-        
-        print(f"🌐 Calling EXTERNAL OCI API (text): {api_url}")
-        
-        # Prepare data for text endpoint
-        data = {
-            'project_id': 'storybook-generator',
-            'subfolder': 'childrens-books',
-            'filename': filename,
-            'content': img_base64  # Base64 encoded image
-        }
-        
-        # Make the API request with timeout
-        response = requests.post(api_url, data=data, timeout=30)
-        
-        print(f"📨 API Response: {response.status_code}")
-        if response.text:
-            print(f"📨 API Response Text: {response.text[:200]}...")
-        
-        if response.status_code == 200:
-            result = response.json()
-            if result['status'] == 'success':
-                return f"✅ {result['message']}"
-            else:
-                return f"❌ API Error: {result.get('message', 'Unknown error')}"
-        else:
-            return f"❌ HTTP Error: {response.status_code} - {response.text}"
-            
-    except requests.exceptions.Timeout:
-        return "❌ API call timed out (30s)"
-    except requests.exceptions.ConnectionError:
-        return "❌ Cannot connect to OCI API"
-    except Exception as e:
-        return f"❌ API upload failed: {str(e)}"
-
-def save_to_oci(image, prompt):
-    """Main function to save image to OCI using best available method"""
-    print("💾 Starting OCI save process...")
-    if DIRECT_OCI_ACCESS:
-        return save_to_oci_direct(image, prompt)
-    else:
-        # Try regular upload first, fallback to text upload if needed
-        result = save_to_oci_via_api(image, prompt)
-        if "❌" in result and "upload-text" not in result:
-            # If regular upload fails, try text endpoint
-            print("🔄 Regular upload failed, trying text endpoint...")
-            return save_to_oci_via_api_text(image, prompt)
-        return result
-
-def generate_storybook_image(prompt):
-    """Generate an image from text prompt and save to OCI"""
+def generate_high_quality_image(prompt, model_choice="dreamshaper-8", style="childrens_book"):
+    """Generate HIGH-QUALITY image from text prompt"""
     try:
         if not prompt or not prompt.strip():
             return None, "❌ Please enter a scene description"
         
-        # Enhance the prompt for better children's book style
-        enhanced_prompt = f"children's book illustration, colorful, whimsical, cute, {prompt}"
+        # Reload model if different choice
+        global pipe
+        if model_choice != "dreamshaper-8":  # Default already loaded
+            pipe = load_model(model_choice)
+        
+        # PROFESSIONAL PROMPT ENHANCEMENT
+        enhanced_prompt, negative_prompt = enhance_prompt(prompt, style)
         
-        print(f"🎨 Generating image for prompt: {enhanced_prompt}")
+        print(f"🎨 Generating HQ image for prompt: {enhanced_prompt}")
+        print(f"🚫 Negative prompt: {negative_prompt[:100]}...")
         
-        # Generate image
+        # Generate HIGH-QUALITY image with better settings
         image = pipe(
-            enhanced_prompt,
-            num_inference_steps=20,
-            guidance_scale=7.5,
-            width=512,
-            height=512
+            prompt=enhanced_prompt,
+            negative_prompt=negative_prompt,
+            num_inference_steps=30,  # More steps = better quality
+            guidance_scale=8.5,      # Higher guidance = better prompt following
+            width=768,               # Higher resolution
+            height=768,
+            generator=torch.Generator(device="cpu").manual_seed(int(time.time()))  # Random seed
         ).images[0]
         
-        print("✅ Image generated successfully!")
+        print("✅ HQ Image generated successfully!")
         
         # Save to OCI
-        print("💾 Saving image to OCI storage...")
-        save_status = save_to_oci(image, prompt)
+        print("💾 Saving HQ image to OCI storage...")
+        save_status = save_to_oci_via_api(image, prompt)
         print(f"Save status: {save_status}")
         
         return image, save_status
         
     except Exception as e:
-        error_msg = f"❌ Generation failed: {str(e)}"
+        error_msg = f"❌ HQ Generation failed: {str(e)}"
         print(error_msg)
         import traceback
         print(f"Traceback: {traceback.format_exc()}")
         return None, error_msg
 
-def check_oci_connection():
-    """Check OCI connection status"""
-    try:
-        if DIRECT_OCI_ACCESS:
-            # Test direct OCI connection
-            namespace = oci_connector.config.get('namespace', 'Not set')
-            bucket = oci_connector.config.get('bucket_name', 'Not set')
-            client_status = "✅ Connected" if oci_connector.client else "❌ Disconnected"
-            
-            return f"""**Direct OCI Connection:**
-- Client: {client_status}
-- Namespace: {namespace}
-- Bucket: {bucket}
-- Initialization Error: {oci_connector.initialization_error or 'None'}"""
-        else:
-            # Test EXTERNAL API connection
-            try:
-                api_url = f"{OCI_API_BASE_URL}/api/health"
-                print(f"🔗 Testing EXTERNAL API: {api_url}")
-                
-                response = requests.get(api_url, timeout=10)
-                if response.status_code == 200:
-                    result = response.json()
-                    return f"✅ EXTERNAL API Connection: Healthy - {result.get('status', 'No status')}"
-                else:
-                    return f"❌ EXTERNAL API Connection: HTTP {response.status_code} - {response.text}"
-            except requests.exceptions.ConnectionError:
-                return "❌ EXTERNAL API Connection: Cannot connect (ConnectionError)"
-            except requests.exceptions.Timeout:
-                return "❌ EXTERNAL API Connection: Timeout after 10 seconds"
-            except Exception as e:
-                return f"❌ EXTERNAL API Connection Error: {str(e)}"
-                
-    except Exception as e:
-        return f"❌ Connection check failed: {str(e)}"
-
-# Create the Gradio interface
-with gr.Blocks(title="Children's Book Illustrator", theme="soft") as demo:
-    gr.Markdown("# 📚 Children's Book Illustrator")
-    gr.Markdown("Generate beautiful storybook images and automatically save them to your OCI storage")
+# Create the Gradio interface with QUALITY OPTIONS
+with gr.Blocks(title="Premium Children's Book Illustrator", theme="soft") as demo:
+    gr.Markdown("# 🎨 Premium Children's Book Illustrator")
+    gr.Markdown("Generate **studio-quality** storybook images with professional results")
     
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=1):
+            gr.Markdown("### 🎯 Quality Settings")
+            
+            model_choice = gr.Dropdown(
+                label="AI Model",
+                choices=list(MODEL_CHOICES.keys()),
+                value="dreamshaper-8",
+                info="Choose the best model for your style"
+            )
+            
+            style_choice = gr.Dropdown(
+                label="Art Style",
+                choices=["childrens_book", "realistic", "fantasy", "anime"],
+                value="childrens_book",
+                info="Select the artistic style"
+            )
+            
             prompt_input = gr.Textbox(
                 label="Scene Description",
-                placeholder="Describe a scene for your storybook...\nExample: A dragon reading a book under a magical tree",
+                placeholder="Describe your scene in detail...\nExample: A friendly dragon reading a giant book under a magical tree with glowing fairies",
                 lines=3
             )
-            generate_btn = gr.Button("🎨 Generate Image", variant="primary")
+            
+            generate_btn = gr.Button("✨ Generate Premium Image", variant="primary")
         
-        with gr.Column():
-            image_output = gr.Image(label="Generated Image", height=400)
+        with gr.Column(scale=2):
+            image_output = gr.Image(label="Generated Image", height=500, show_download_button=True)
             status_output = gr.Textbox(label="Status", interactive=False, lines=3)
     
-    # Debug section
-    with gr.Accordion("🔧 Debug Information", open=False):
-        gr.Markdown(f"""
-        **Current Configuration:**
-        - OCI Access: **{'Direct' if DIRECT_OCI_ACCESS else 'External API'}**
-        - API URL: {OCI_API_BASE_URL}
-        - Model: Stable Diffusion v1.5
-        - Storage Path: `storybook-generator/childrens-books/`
-        """)
+    # Examples section
+    with gr.Accordion("💡 Prompt Examples & Tips", open=False):
+        gr.Markdown("""
+        ## 🎨 Professional Prompt Examples:
+        
+        **Best Results:**
+        - "A cute baby dragon learning to read with an old wise turtle in a magical forest"
+        - "Group of animal friends having a picnic under a giant glowing mushroom"
+        - "Little girl exploring a castle made of candy with talking animals"
+        - "Space adventure with friendly robots and colorful aliens"
         
-        debug_btn = gr.Button("🔄 Check OCI Connection", variant="secondary")
-        debug_output = gr.Textbox(label="Debug Info", interactive=False, lines=8)
+        ## 🚫 Avoid:
+        - "dragon" → **"friendly cartoon dragon reading book"**
+        - "cat" → **"cute kitten playing with yarn ball"**
+        - "tree" → **"magical tree with glowing fairies"**
+        
+        ## ⚡ Pro Tips:
+        1. **Be descriptive** - Add colors, emotions, settings
+        2. **Specify style** - "watercolor", "cartoon", "realistic"
+        3. **Add details** - "with glowing effects", "in a magical forest"
+        4. **Use positive language** - "happy", "friendly", "colorful"
+        """)
     
+    # Debug section
+    with gr.Accordion("🔧 Advanced Settings", open=False):
+        debug_btn = gr.Button("🔄 Check System Status", variant="secondary")
+        debug_output = gr.Textbox(label="System Info", interactive=False, lines=4)
+    
+    def check_system_status():
+        """Check system status"""
+        return f"""**System Status:**
+- Model: {model_choice.value}
+- Style: {style_choice.value}
+- OCI API: {OCI_API_BASE_URL}
+- Ready for premium image generation!"""
+
     # Connect buttons to functions
     generate_btn.click(
-        fn=generate_storybook_image,
-        inputs=prompt_input,
+        fn=generate_high_quality_image,
+        inputs=[prompt_input, model_choice, style_choice],
         outputs=[image_output, status_output]
     )
     
     debug_btn.click(
-        fn=check_oci_connection,
+        fn=check_system_status,
         inputs=None,
         outputs=debug_output
     )
@@ -347,7 +299,8 @@ def get_app():
     return demo
 
 if __name__ == "__main__":
-    print("🚀 Starting Children's Book Illustrator...")
-    print(f"📦 OCI Access: {'Direct' if DIRECT_OCI_ACCESS else 'External API'}")
-    print(f"🌐 API URL: {OCI_API_BASE_URL}")
+    print("🚀 Starting Premium Children's Book Illustrator...")
+    print("🎨 Loaded with HIGH-QUALITY models:")
+    for model in MODEL_CHOICES:
+        print(f"   - {model}")
     demo.launch(server_name="0.0.0.0", server_port=7860)