Update app.py

app.py

@@ -3,130 +3,142 @@ import numpy as np
 import random
 import torch
 import gc
-from typing import Optional, Tuple
 import warnings
+from typing import Optional, Tuple
+import os
+
+# Suppress warnings
 warnings.filterwarnings("ignore")
+os.environ["TRANSFORMERS_VERBOSITY"] = "error"
 
-# CPU-optimized imports
-from diffusers import (
-    StableDiffusionPipeline, 
-    DiffusionPipeline,
-    DDIMScheduler,
-    DPMSolverMultistepScheduler
-)
-from transformers import CLIPTokenizer
+# Import compatible with older versions
+try:
+    from diffusers import StableDiffusionPipeline, DiffusionPipeline
+    from diffusers import DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler
+    DIFFUSERS_AVAILABLE = True
+except ImportError:
+    DIFFUSERS_AVAILABLE = False
 
-# Force CPU usage and optimize for low-resource environment
+# Fallback to transformers + torch if diffusers fails
+if not DIFFUSERS_AVAILABLE:
+    try:
+        from transformers import CLIPTokenizer, CLIPTextModel
+        from transformers import AutoTokenizer, AutoModel
+        TRANSFORMERS_AVAILABLE = True
+    except ImportError:
+        TRANSFORMERS_AVAILABLE = False
+
+# Force CPU and optimize
 device = "cpu"
-torch.set_num_threads(2)  # Match vCPU count
+torch.set_num_threads(2)
 MAX_SEED = np.iinfo(np.int32).max
 
-# CPU-optimized model configurations
+# Compatible model configurations for older versions
 MODEL_CONFIGS = {
-    "🚀 Tiny SD (Ultra Fast)": {
-        "repo_id": "nota-ai/bk-sdm-small",
-        "torch_dtype": torch.float32,
-        "max_resolution": 512,
-        "default_steps": 10,
-        "default_guidance": 6.0,
-        "memory_usage": "Very Low",
-        "speed": "Ultra Fast"
-    },
-    "⚡ SD 1.4 (Fast)": {
+    "🚀 CompVis SD 1.4 (Fast)": {
         "repo_id": "CompVis/stable-diffusion-v1-4",
+        "revision": "main",
         "torch_dtype": torch.float32,
         "max_resolution": 512,
         "default_steps": 15,
         "default_guidance": 7.5,
         "memory_usage": "Low",
-        "speed": "Fast"
+        "compatible": True
     },
-    "🎨 SD 1.5 (Balanced)": {
+    "⚡ Runway SD 1.5 (Recommended)": {
         "repo_id": "runwayml/stable-diffusion-v1-5",
+        "revision": "main", 
         "torch_dtype": torch.float32,
         "max_resolution": 512,
         "default_steps": 20,
         "default_guidance": 7.5,
         "memory_usage": "Medium",
-        "speed": "Medium"
+        "compatible": True
     },
-    "🖼️ OpenJourney (Artistic)": {
-        "repo_id": "prompthero/openjourney",
+    "🎨 OpenJourney v4 (Artistic)": {
+        "repo_id": "prompthero/openjourney-v4",
+        "revision": "main",
         "torch_dtype": torch.float32,
         "max_resolution": 512,
         "default_steps": 18,
         "default_guidance": 8.0,
         "memory_usage": "Medium",
-        "speed": "Medium"
+        "compatible": True
     },
-    "🌟 Dreamlike (Quality)": {
-        "repo_id": "dreamlike-art/dreamlike-diffusion-1.0",
+    "🌟 Anything v3 (Anime Style)": {
+        "repo_id": "Linaqruf/anything-v3.0",
+        "revision": "main",
         "torch_dtype": torch.float32,
-        "max_resolution": 448,
-        "default_steps": 25,
-        "default_guidance": 8.0,
-        "memory_usage": "Medium-High",
-        "speed": "Slower"
+        "max_resolution": 512,
+        "default_steps": 20,
+        "default_guidance": 8.5,
+        "memory_usage": "Medium",
+        "compatible": True
     }
 }
 
-# Global variables for memory management
+# Global pipeline cache
 current_pipeline = None
 current_model_name = None
 
 def clear_memory():
-    """Aggressive memory cleanup for CPU environment"""
+    """Aggressive memory cleanup"""
     global current_pipeline
     if current_pipeline is not None:
         del current_pipeline
         current_pipeline = None
-    
-    # Force garbage collection
     gc.collect()
-    
-    # Clear any cached models
-    torch.cuda.empty_cache() if torch.cuda.is_available() else None
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
 
-def load_model_cpu_optimized(model_name: str):
-    """Load model with maximum CPU optimization"""
+def load_pipeline_safe(model_name: str):
+    """Load pipeline with maximum compatibility"""
     global current_pipeline, current_model_name
     
-    # Return cached pipeline if same model
     if current_model_name == model_name and current_pipeline is not None:
         return current_pipeline, "✅ Using cached model"
     
-    # Clear previous model
     clear_memory()
     
+    if not DIFFUSERS_AVAILABLE:
+        return None, "❌ Diffusers library not available. Please install compatible versions."
+    
     try:
         config = MODEL_CONFIGS[model_name]
         
-        # Load with CPU optimizations
+        # Load with maximum compatibility
         pipe = StableDiffusionPipeline.from_pretrained(
             config["repo_id"],
+            revision=config.get("revision", "main"),
             torch_dtype=config["torch_dtype"],
-            safety_checker=None,  # Disable for speed
+            safety_checker=None,
             requires_safety_checker=False,
-            use_safetensors=False,  # Faster loading on CPU
-            local_files_only=False
+            use_auth_token=False,
+            cache_dir=None,
+            local_files_only=False,
+            low_cpu_mem_usage=True,
+            ignore_mismatched_sizes=True
         )
         
-        # Apply CPU-specific optimizations
+        # Move to CPU and optimize
         pipe = pipe.to(device)
         
         # Enable attention slicing for memory efficiency
-        pipe.enable_attention_slicing(1)
+        if hasattr(pipe, 'enable_attention_slicing'):
+            pipe.enable_attention_slicing(1)
         
-        # Use memory efficient scheduler
-        pipe.scheduler = DPMSolverMultistepScheduler.from_config(
-            pipe.scheduler.config,
-            use_karras_sigmas=True,
-            algorithm_type="dpmsolver++"
-        )
+        # Use compatible scheduler
+        if hasattr(pipe, 'scheduler'):
+            try:
+                pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+            except:
+                pass  # Keep original scheduler if DDIMScheduler fails
         
-        # Additional CPU optimizations
-        pipe.unet.to(memory_format=torch.channels_last)
-        pipe.vae.to(memory_format=torch.channels_last)
+        # Memory optimizations
+        if hasattr(pipe.unet, 'to'):
+            pipe.unet.to(memory_format=torch.channels_last)
+        if hasattr(pipe.vae, 'to'):
+            pipe.vae.to(memory_format=torch.channels_last)
         
         current_pipeline = pipe
         current_model_name = model_name
@@ -134,9 +146,15 @@ def load_model_cpu_optimized(model_name: str):
         return pipe, f"✅ {model_name} loaded successfully!"
         
     except Exception as e:
-        return None, f"❌ Failed to load {model_name}: {str(e)}"
+        error_msg = str(e)
+        if "torch" in error_msg and "float8" in error_msg:
+            return None, "❌ Version incompatibility. Please update PyTorch to >= 2.1 or use older Diffusers version."
+        elif "out of memory" in error_msg.lower():
+            return None, "❌ Out of memory. Try using a different model or restart the space."
+        else:
+            return None, f"❌ Failed to load model: {error_msg[:200]}..."
 
-def generate_image_cpu(
+def generate_image_compatible(
     model_name: str,
     prompt: str,
     negative_prompt: str,
@@ -152,8 +170,8 @@ def generate_image_cpu(
     if not prompt.strip():
         return None, seed, "⚠️ Please enter a prompt"
     
-    # Load model
-    pipe, status = load_model_cpu_optimized(model_name)
+    # Load pipeline
+    pipe, status = load_pipeline_safe(model_name)
     if pipe is None:
         return None, seed, status
     
@@ -164,7 +182,7 @@ def generate_image_cpu(
         
         generator = torch.Generator().manual_seed(seed)
         
-        # CPU-specific constraints
+        # Apply constraints for CPU
         config = MODEL_CONFIGS[model_name]
         max_res = config["max_resolution"]
         width = min(width, max_res)
@@ -174,27 +192,42 @@ def generate_image_cpu(
         width = (width // 8) * 8
         height = (height // 8) * 8
         
-        # Limit steps for CPU
+        # Limit steps for CPU performance
         num_inference_steps = min(num_inference_steps, 30)
         
-        progress(0, desc="Starting generation...")
+        progress(0, desc="Initializing generation...")
         
-        # Generate with CPU optimizations
+        # Generate with error handling
         with torch.no_grad():
-            result = pipe(
-                prompt=prompt,
-                negative_prompt=negative_prompt if negative_prompt.strip() else None,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_inference_steps,
-                width=width,
-                height=height,
-                generator=generator,
-                callback_on_step_end=lambda step, timestep, latents: progress(step/num_inference_steps)
-            )
-        
-        image = result.images[0]
+            try:
+                result = pipe(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt if negative_prompt.strip() else None,
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=num_inference_steps,
+                    width=width,
+                    height=height,
+                    generator=generator,
+                )
+                image = result.images[0]
+            except Exception as gen_error:
+                # Fallback: try with minimal parameters
+                if "memory" in str(gen_error).lower():
+                    width, height = 256, 256
+                    num_inference_steps = 10
+                    result = pipe(
+                        prompt=prompt,
+                        guidance_scale=guidance_scale,
+                        num_inference_steps=num_inference_steps,
+                        width=width,
+                        height=height,
+                        generator=generator,
+                    )
+                    image = result.images[0]
+                else:
+                    raise gen_error
         
-        # Memory cleanup after generation
+        # Cleanup
         del result
         gc.collect()
         
@@ -203,220 +236,213 @@ def generate_image_cpu(
     except Exception as e:
         error_msg = f"❌ Generation failed: {str(e)}"
         if "memory" in str(e).lower() or "out of" in str(e).lower():
-            error_msg += "\n💡 Try: Smaller resolution (256x256), fewer steps (10-15), or Tiny SD model"
+            error_msg += "\n💡 Try: 256x256 resolution, 10 steps, or restart the space"
+        elif "CUDA" in str(e):
+            error_msg += "\n💡 CUDA error detected, using CPU fallback"
         return None, seed, error_msg
 
-# Optimized example prompts for CPU generation
+# Check system compatibility
+def check_system():
+    """Check system compatibility and return status"""
+    status = []
+    
+    # Check PyTorch
+    torch_version = torch.__version__
+    status.append(f"🔧 PyTorch: {torch_version}")
+    
+    # Check diffusers
+    if DIFFUSERS_AVAILABLE:
+        try:
+            import diffusers
+            status.append(f"✅ Diffusers: {diffusers.__version__}")
+        except:
+            status.append("⚠️ Diffusers: Version unknown")
+    else:
+        status.append("❌ Diffusers: Not available")
+    
+    # Check transformers
+    if TRANSFORMERS_AVAILABLE:
+        try:
+            import transformers
+            status.append(f"✅ Transformers: {transformers.__version__}")
+        except:
+            status.append("⚠️ Transformers: Version unknown")
+    else:
+        status.append("❌ Transformers: Not available")
+    
+    # Check device
+    status.append(f"🖥️ Device: {device.upper()}")
+    status.append(f"🧵 CPU Threads: {torch.get_num_threads()}")
+    
+    return "\n".join(status)
+
+# Example prompts optimized for compatibility
 examples = [
-    "a cute cat sitting in a garden, digital art",
-    "mountain landscape at sunset, beautiful painting",
-    "portrait of a smiling person, photography",
-    "colorful flowers in a vase, oil painting",
-    "futuristic city skyline, concept art",
-    "peaceful forest path, nature photography"
+    "a beautiful landscape with mountains and lake",
+    "portrait of a cat, digital art style", 
+    "colorful flowers in a garden, painting",
+    "medieval castle on a hill, fantasy art",
+    "astronaut in space, realistic style",
+    "cozy coffee shop interior, warm lighting"
 ]
 
-# CSS optimized for CPU performance (minimal animations)
+# Minimal CSS for compatibility
 css = """
 #col-container {
     margin: 0 auto;
     max-width: 800px;
     padding: 20px;
 }
-.model-card {
+.info-box {
     padding: 15px;
     margin: 10px 0;
     border-radius: 8px;
-    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
-    border: 1px solid #ddd;
-}
-.cpu-optimized {
-    background: linear-gradient(135deg, #a8edea 0%, #fed6e3 100%);
-    font-weight: bold;
+    background-color: #f8f9fa;
+    border-left: 4px solid #007bff;
 }
-.status-text {
+.model-info {
     padding: 10px;
+    margin: 5px 0;
     border-radius: 5px;
-    margin: 10px 0;
+    background-color: #e9ecef;
+    font-size: 0.9em;
 }
-.status-success { background-color: #d4edda; color: #155724; }
-.status-error { background-color: #f8d7da; color: #721c24; }
-.status-warning { background-color: #fff3cd; color: #856404; }
+.status-success { background-color: #d4edda; color: #155724; padding: 10px; border-radius: 5px; }
+.status-error { background-color: #f8d7da; color: #721c24; padding: 10px; border-radius: 5px; }
 """
 
-# Main interface optimized for CPU
-with gr.Blocks(css=css, title="CPU-Optimized AI Image Generator") as demo:
+# Main interface with maximum compatibility
+with gr.Blocks(css=css, title="Compatible AI Image Generator") as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown("""
-        # 🖥️ CPU-Optimized AI Image Generator
-        ### Specially optimized for CPU Basic (2 vCPUs, 16GB RAM)
-        """)
+        gr.Markdown("# 🎨 Compatible AI Image Generator")
+        gr.Markdown("### Optimized for CPU Basic with older PyTorch versions")
+        
+        # System status
+        with gr.Accordion("🔧 System Status", open=False):
+            system_status = gr.Markdown(check_system())
         
-        # System info
-        gr.Markdown("""
-        <div class="model-card cpu-optimized">
-        📊 <strong>System Configuration:</strong> CPU Basic | 2 vCPUs | 16GB RAM | PyTorch CPU
-        </div>
-        """)
+        # Warning for incompatible systems
+        if not DIFFUSERS_AVAILABLE:
+            gr.Markdown("""
+            <div style="background-color: #f8d7da; color: #721c24; padding: 15px; border-radius: 8px; margin: 10px 0;">
+            ⚠️ <strong>Compatibility Issue Detected</strong><br>
+            Please update your requirements.txt with compatible versions:<br>
+            <code>
+            torch>=2.1.0<br>
+            diffusers>=0.21.0,<0.25.0<br>
+            transformers>=4.25.0,<4.35.0<br>
+            accelerate>=0.20.0<br>
+            </code>
+            </div>
+            """)
         
+        # Main interface
         with gr.Row():
             with gr.Column(scale=3):
                 prompt = gr.Textbox(
-                    label="✨ Your Creative Prompt",
-                    placeholder="Describe the image you want to create...",
-                    lines=3,
-                    max_lines=5
+                    label="✨ Describe your image",
+                    placeholder="Enter your creative prompt here...",
+                    lines=3
                 )
             with gr.Column(scale=1):
                 generate_btn = gr.Button(
-                    "🎨 Generate Image", 
+                    "🎨 Generate", 
                     variant="primary", 
-                    size="lg",
-                    elem_classes="generate-button"
+                    size="lg"
                 )
         
-        # Model selection with detailed info
-        with gr.Row():
-            model_dropdown = gr.Dropdown(
-                choices=list(MODEL_CONFIGS.keys()),
-                value="🚀 Tiny SD (Ultra Fast)",
-                label="🤖 AI Model Selection",
-                info="Choose based on speed vs quality preference"
-            )
-        
-        # Model info display
-        model_info_display = gr.Markdown("", elem_classes="model-card")
-        
-        # Generated image display
-        result_image = gr.Image(
-            label="🖼️ Generated Image", 
-            height=400,
-            show_label=True
+        # Model selection
+        model_dropdown = gr.Dropdown(
+            choices=list(MODEL_CONFIGS.keys()),
+            value="⚡ Runway SD 1.5 (Recommended)",
+            label="🤖 AI Model",
+            info="All models optimized for CPU usage"
         )
         
-        # Status display
-        status_display = gr.Markdown(
-            "🚀 Ready to generate! Select a model and enter your prompt.", 
-            elem_classes="status-text status-success"
-        )
+        # Model info
+        model_info = gr.Markdown("", elem_classes="model-info")
         
-        # Advanced settings in accordion
-        with gr.Accordion("⚙️ Advanced Settings", open=False):
-            with gr.Column():
-                negative_prompt = gr.Textbox(
-                    label="🚫 Negative Prompt (Optional)",
-                    placeholder="What you don't want in the image...",
-                    lines=2
-                )
-                
-                with gr.Row():
-                    seed = gr.Slider(0, MAX_SEED, value=0, label="🎲 Seed")
-                    randomize_seed = gr.Checkbox(label="🔄 Random Seed", value=True)
-                
-                with gr.Row():
-                    width = gr.Slider(
-                        256, 512, value=384, step=64, 
-                        label="📏 Width",
-                        info="Lower = faster generation"
-                    )
-                    height = gr.Slider(
-                        256, 512, value=384, step=64, 
-                        label="📐 Height",
-                        info="Lower = faster generation"
-                    )
-                
-                with gr.Row():
-                    guidance_scale = gr.Slider(
-                        1.0, 15.0, value=7.5, step=0.5,
-                        label="🎯 Guidance Scale",
-                        info="How closely to follow the prompt"
-                    )
-                    num_inference_steps = gr.Slider(
-                        5, 30, value=15, step=1,
-                        label="🔄 Steps",
-                        info="More steps = better quality but slower"
-                    )
+        # Result
+        result_image = gr.Image(label="Generated Image", height=400)
+        status_text = gr.Markdown("🚀 Ready to generate!", elem_classes="status-success")
         
-        # CPU Performance Tips
-        with gr.Accordion("💡 CPU Optimization Tips", open=False):
-            gr.Markdown("""
-            ### 🚀 For Fastest Generation:
-            - Use **Tiny SD** model
-            - Set resolution to **256x256** or **384x384**
-            - Use **10-15 steps**
-            - Keep guidance scale around **6-8**
+        # Settings
+        with gr.Accordion("⚙️ Generation Settings", open=False):
+            negative_prompt = gr.Textbox(
+                label="🚫 Negative Prompt",
+                placeholder="What you don't want...",
+                lines=2
+            )
+            
+            with gr.Row():
+                seed = gr.Slider(0, MAX_SEED, value=0, label="🎲 Seed")
+                randomize_seed = gr.Checkbox(label="🔄 Random", value=True)
             
-            ### 🎨 For Best Quality:
-            - Use **Dreamlike** or **SD 1.5** model
-            - Set resolution to **512x512** (max)
-            - Use **20-25 steps**
-            - Guidance scale **7-9**
+            with gr.Row():
+                width = gr.Slider(256, 512, value=384, step=64, label="Width")
+                height = gr.Slider(256, 512, value=384, step=64, label="Height")
             
-            ### ⚡ Memory Saving:
-            - Generate one image at a time
-            - Use shorter prompts when possible
-            - Avoid very high guidance scales (>12)
+            with gr.Row():
+                guidance_scale = gr.Slider(1.0, 15.0, value=7.5, step=0.5, label="Guidance")
+                num_inference_steps = gr.Slider(5, 30, value=15, step=1, label="Steps")
+        
+        # Tips
+        with gr.Accordion("💡 Performance Tips", open=False):
+            gr.Markdown("""
+            ### For Best Results on CPU Basic:
+            - **Fast Generation**: Use 256x256, 10-15 steps
+            - **Quality Generation**: Use 384x384, 20 steps  
+            - **Maximum Quality**: Use 512x512, 25 steps (slower)
+            - **Memory Issues**: Restart the space if you get memory errors
+            - **Compatibility**: Update PyTorch to 2.1+ for best performance
             """)
         
         # Examples
-        gr.Examples(
-            examples=examples,
-            inputs=[prompt],
-            label="💡 Example Prompts (Click to try!)"
-        )
+        gr.Examples(examples=examples, inputs=[prompt])
         
-        # Footer with helpful info
-        gr.Markdown("""
-        ---
-        <div style="text-align: center; color: #666; font-size: 0.9em;">
-        🖥️ Optimized for CPU Basic | Generation time: 30s-3min depending on settings
-        </div>
-        """)
+        # Footer
+        gr.Markdown("---")
+        gr.Markdown("🖥️ **CPU Optimized** | Generation time: 30s-3min depending on settings")
     
-    # Function to update model info
+    # Update model info
     def update_model_info(model_name):
-        config = MODEL_CONFIGS[model_name]
-        info = f"""
-        <div class="model-card">
-        <strong>{model_name}</strong><br>
-        📊 <strong>Memory Usage:</strong> {config['memory_usage']} | 
-        ⚡ <strong>Speed:</strong> {config['speed']}<br>
-        📏 <strong>Max Resolution:</strong> {config['max_resolution']}px | 
-        🔄 <strong>Recommended Steps:</strong> {config['default_steps']}<br>
-        🎯 <strong>Recommended Guidance:</strong> {config['default_guidance']}
-        </div>
-        """
-        return info, config['default_steps'], config['default_guidance']
+        if model_name in MODEL_CONFIGS:
+            config = MODEL_CONFIGS[model_name]
+            info = f"""
+            **Memory Usage:** {config['memory_usage']} | **Max Resolution:** {config['max_resolution']}px  
+            **Recommended:** {config['default_steps']} steps, {config['default_guidance']} guidance
+            """
+            return info, config['default_steps'], config['default_guidance']
+        return "", 15, 7.5
     
     # Event handlers
     model_dropdown.change(
         update_model_info,
         inputs=[model_dropdown],
-        outputs=[model_info_display, num_inference_steps, guidance_scale]
+        outputs=[model_info, num_inference_steps, guidance_scale]
     )
     
-    # Generation handler
     generate_btn.click(
-        generate_image_cpu,
+        generate_image_compatible,
         inputs=[
             model_dropdown, prompt, negative_prompt,
             seed, randomize_seed, width, height,
             guidance_scale, num_inference_steps
         ],
-        outputs=[result_image, seed, status_display]
+        outputs=[result_image, seed, status_text]
     )
     
-    # Auto-trigger model info update on load
+    # Initialize model info
     demo.load(
         update_model_info,
         inputs=[model_dropdown],
-        outputs=[model_info_display, num_inference_steps, guidance_scale]
+        outputs=[model_info, num_inference_steps, guidance_scale]
     )
 
 if __name__ == "__main__":
     demo.launch(
         share=True,
-        server_name="0.0.0.0",
+        server_name="0.0.0.0",  
         server_port=7860,
         show_error=True,
         quiet=True