Update Dockerfile

Dockerfile

@@ -1,30 +1,471 @@
-# Base image llama.cpp se hi le rahe hain
-FROM ghcr.io/ggml-org/llama.cpp:full
-
-WORKDIR /app
-
-# System dependencies install karna
-RUN apt update && apt install -y python3 python3-pip python3-venv
-
-# Virtual environment setup
-RUN python3 -m venv /opt/venv
-ENV PATH="/opt/venv/bin:$PATH"
-
-# Python tools install karna
-RUN pip install -U pip huggingface_hub
-
-# Model download karna (Qwen2.5-Coder-3B-Instruct)
-RUN python3 -c 'from huggingface_hub import hf_hub_download; \
-    repo="bartowski/Qwen2.5-Coder-3B-Instruct-GGUF"; \
-    hf_hub_download(repo_id=repo, filename="Qwen2.5-Coder-3B-Instruct-Q4_K_M.gguf", local_dir="/app")'
-
-# Server start karne ka command
-CMD ["--server", \
-     "-m", "/app/Qwen2.5-Coder-3B-Instruct-Q4_K_M.gguf", \
-     "--host", "0.0.0.0", \
-     "--port", "7860", \
-     "-t", "2", \
-     "--cache-type-k", "q8_0", \
-     "--cache-type-v", "iq4_nl", \
-     "-c", "18192", \
-     "-n", "4096"]
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+import os
+from pathlib import Path
+import logging
+from typing import Optional, Tuple
+import time
+from diffusers import StableDiffusionXLImg2ImgPipeline, StableDiffusionXLPipeline
+import gguf
+import ctypes
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# ===================== DEVICE & CONFIG =====================
+DEVICE = "cpu"  # CPU-only for HuggingFace Spaces
+DTYPE = torch.float32  # CPU ke liye float32 use karna zaroori hai
+
+# FLUX.1 Model Config
+MODEL_CONFIG = {
+    "model_id": "unsloth/FLUX.1-Kontext-dev-GGUF",
+    "gguf_file": "flux1-dev-Q4_K_M.gguf",  # Quantized version
+    "use_safetensors": False,
+}
+
+# Fallback to Stability if GGUF load nahi hota
+FALLBACK_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
+
+# ===================== GLOBAL STATE =====================
+class PipelineManager:
+    def __init__(self):
+        self.txt2img_pipe = None
+        self.img2img_pipe = None
+        self.model_loaded = False
+        self.current_model = None
+        self.load_lock = False
+        
+    def load_flux_gguf(self):
+        """Load FLUX.1 GGUF model from HF"""
+        try:
+            logger.info("📥 Loading FLUX.1 GGUF model...")
+            
+            # Try loading from HF hub with GGUF support
+            from huggingface_hub import hf_hub_download
+            
+            model_path = hf_hub_download(
+                repo_id=MODEL_CONFIG["model_id"],
+                filename=MODEL_CONFIG["gguf_file"],
+                cache_dir="./models"
+            )
+            
+            logger.info(f"✅ GGUF downloaded: {model_path}")
+            
+            # Load using llama-cpp for GGUF
+            from llama_cpp import Llama
+            llm = Llama.from_pretrained(
+                MODEL_CONFIG["model_id"],
+                filename=MODEL_CONFIG["gguf_file"],
+                n_gpu_layers=-1,  # CPU-only, to GPU layers = -1
+                n_threads=os.cpu_count(),
+                verbose=False
+            )
+            
+            logger.info("✅ FLUX.1 GGUF loaded successfully!")
+            self.current_model = "flux1-gguf"
+            return True
+            
+        except Exception as e:
+            logger.warning(f"⚠️ GGUF loading failed: {e}")
+            return False
+    
+    def load_flux_diffusers(self):
+        """Load FLUX.1 using diffusers (as fallback)"""
+        try:
+            logger.info("📥 Loading FLUX.1 via diffusers...")
+            
+            # Lighter quantized version for CPU
+            self.txt2img_pipe = StableDiffusionXLPipeline.from_pretrained(
+                "stabilityai/stable-diffusion-xl-base-1.0",
+                torch_dtype=DTYPE,
+                use_safetensors=True,
+                variant="fp32"
+            )
+            
+            self.txt2img_pipe = self.txt2img_pipe.to(DEVICE)
+            self.txt2img_pipe.enable_attention_slicing()  # CPU ke liye memory optimize
+            
+            self.img2img_pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+                "stabilityai/stable-diffusion-xl-refiner-1.0",
+                torch_dtype=DTYPE,
+                use_safetensors=True,
+                variant="fp32"
+            )
+            
+            self.img2img_pipe = self.img2img_pipe.to(DEVICE)
+            self.img2img_pipe.enable_attention_slicing()
+            
+            logger.info("✅ FLUX.1 diffusers pipeline loaded!")
+            self.current_model = "flux1-diffusers"
+            self.model_loaded = True
+            return True
+            
+        except Exception as e:
+            logger.error(f"❌ Diffusers loading failed: {e}")
+            return False
+    
+    def initialize(self):
+        """Initialize pipelines"""
+        if self.load_lock:
+            return
+        
+        self.load_lock = True
+        
+        # Try GGUF first
+        if not self.load_flux_gguf():
+            # Fallback to diffusers
+            self.load_flux_diffusers()
+        
+        self.load_lock = False
+    
+    def generate_txt2img(
+        self,
+        prompt: str,
+        negative_prompt: str = "",
+        num_inference_steps: int = 20,
+        guidance_scale: float = 7.5,
+        height: int = 768,
+        width: int = 768,
+        seed: int = -1
+    ) -> Image.Image:
+        """Text-to-Image generation"""
+        
+        if not self.model_loaded:
+            raise RuntimeError("Model not initialized")
+        
+        if seed == -1:
+            seed = int(time.time())
+        
+        generator = torch.Generator(device=DEVICE).manual_seed(seed)
+        
+        logger.info(f"🎨 Generating txt2img: {prompt[:50]}...")
+        
+        with torch.no_grad():  # CPU memory optimization
+            image = self.txt2img_pipe(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                height=height,
+                width=width,
+                generator=generator
+            ).images[0]
+        
+        logger.info("✅ Generation complete!")
+        return image
+    
+    def generate_img2img(
+        self,
+        prompt: str,
+        image: Image.Image,
+        negative_prompt: str = "",
+        num_inference_steps: int = 20,
+        guidance_scale: float = 7.5,
+        strength: float = 0.8,
+        seed: int = -1
+    ) -> Image.Image:
+        """Image-to-Image generation"""
+        
+        if not self.model_loaded:
+            raise RuntimeError("Model not initialized")
+        
+        if seed == -1:
+            seed = int(time.time())
+        
+        generator = torch.Generator(device=DEVICE).manual_seed(seed)
+        
+        # Resize image to 768x768
+        image = image.resize((768, 768), Image.Resampling.LANCZOS)
+        
+        logger.info(f"🖼️ Generating img2img: {prompt[:50]}...")
+        
+        with torch.no_grad():
+            image = self.img2img_pipe(
+                prompt=prompt,
+                image=image,
+                negative_prompt=negative_prompt,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                strength=strength,
+                generator=generator
+            ).images[0]
+        
+        logger.info("✅ Generation complete!")
+        return image
+
+# Initialize global pipeline manager
+pipeline_manager = PipelineManager()
+
+# ===================== GRADIO INTERFACE =====================
+
+def generate_txt2img_ui(
+    prompt: str,
+    negative_prompt: str,
+    num_steps: int,
+    guidance: float,
+    height: int,
+    width: int,
+    seed: int
+):
+    """Gradio wrapper for txt2img"""
+    try:
+        if not pipeline_manager.model_loaded:
+            return None, "❌ Model not loaded yet. Please wait..."
+        
+        image = pipeline_manager.generate_txt2img(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_inference_steps=num_steps,
+            guidance_scale=guidance,
+            height=height,
+            width=width,
+            seed=seed
+        )
+        
+        return image, "✅ Generation successful!"
+        
+    except Exception as e:
+        logger.error(f"Error in txt2img: {e}")
+        return None, f"❌ Error: {str(e)}"
+
+def generate_img2img_ui(
+    prompt: str,
+    image: Image.Image,
+    negative_prompt: str,
+    num_steps: int,
+    guidance: float,
+    strength: float,
+    seed: int
+):
+    """Gradio wrapper for img2img"""
+    try:
+        if image is None:
+            return None, "❌ Please upload an image first"
+        
+        if not pipeline_manager.model_loaded:
+            return None, "❌ Model not loaded yet. Please wait..."
+        
+        result = pipeline_manager.generate_img2img(
+            prompt=prompt,
+            image=image,
+            negative_prompt=negative_prompt,
+            num_inference_steps=num_steps,
+            guidance_scale=guidance,
+            strength=strength,
+            seed=seed
+        )
+        
+        return result, "✅ Generation successful!"
+        
+    except Exception as e:
+        logger.error(f"Error in img2img: {e}")
+        return None, f"❌ Error: {str(e)}"
+
+# ===================== GRADIO BLOCKS UI =====================
+
+with gr.Blocks(title="FLUX.1 GGUF - Text/Image Generator", theme=gr.themes.Soft()) as demo:
+    
+    gr.Markdown("# 🎨 FLUX.1 GGUF Generator")
+    gr.Markdown("CPU-optimized Text-to-Image & Image-to-Image generation")
+    
+    # Status indicator
+    status = gr.Textbox(value="⏳ Loading model...", interactive=False, label="Status")
+    
+    with gr.Tabs():
+        
+        # ===== TAB 1: TEXT-TO-IMAGE =====
+        with gr.Tab("📝 Text-to-Image"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    txt2img_prompt = gr.Textbox(
+                        label="Prompt",
+                        placeholder="A beautiful sunset over mountains...",
+                        lines=3
+                    )
+                    txt2img_negative = gr.Textbox(
+                        label="Negative Prompt",
+                        placeholder="blurry, low quality, distorted...",
+                        lines=2
+                    )
+                    
+                    with gr.Row():
+                        txt2img_height = gr.Slider(
+                            minimum=256,
+                            maximum=1024,
+                            value=768,
+                            step=64,
+                            label="Height"
+                        )
+                        txt2img_width = gr.Slider(
+                            minimum=256,
+                            maximum=1024,
+                            value=768,
+                            step=64,
+                            label="Width"
+                        )
+                    
+                    with gr.Row():
+                        txt2img_steps = gr.Slider(
+                            minimum=1,
+                            maximum=50,
+                            value=20,
+                            step=1,
+                            label="Inference Steps"
+                        )
+                        txt2img_guidance = gr.Slider(
+                            minimum=1,
+                            maximum=15,
+                            value=7.5,
+                            step=0.5,
+                            label="Guidance Scale"
+                        )
+                    
+                    txt2img_seed = gr.Number(
+                        value=-1,
+                        label="Seed (-1 = random)",
+                        precision=0
+                    )
+                    
+                    txt2img_generate = gr.Button(
+                        "🎨 Generate Image",
+                        variant="primary",
+                        size="lg"
+                    )
+                
+                with gr.Column(scale=1):
+                    txt2img_output = gr.Image(label="Generated Image", type="pil")
+                    txt2img_status = gr.Textbox(interactive=False, label="Result")
+            
+            txt2img_generate.click(
+                generate_txt2img_ui,
+                inputs=[
+                    txt2img_prompt,
+                    txt2img_negative,
+                    txt2img_steps,
+                    txt2img_guidance,
+                    txt2img_height,
+                    txt2img_width,
+                    txt2img_seed
+                ],
+                outputs=[txt2img_output, txt2img_status]
+            )
+        
+        # ===== TAB 2: IMAGE-TO-IMAGE =====
+        with gr.Tab("🖼️ Image-to-Image"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    img2img_input = gr.Image(
+                        label="Input Image",
+                        type="pil"
+                    )
+                    
+                    img2img_prompt = gr.Textbox(
+                        label="Prompt",
+                        placeholder="Transform the image to...",
+                        lines=3
+                    )
+                    img2img_negative = gr.Textbox(
+                        label="Negative Prompt",
+                        placeholder="blurry, low quality...",
+                        lines=2
+                    )
+                    
+                    with gr.Row():
+                        img2img_steps = gr.Slider(
+                            minimum=1,
+                            maximum=50,
+                            value=20,
+                            step=1,
+                            label="Inference Steps"
+                        )
+                        img2img_guidance = gr.Slider(
+                            minimum=1,
+                            maximum=15,
+                            value=7.5,
+                            step=0.5,
+                            label="Guidance Scale"
+                        )
+                    
+                    img2img_strength = gr.Slider(
+                        minimum=0,
+                        maximum=1,
+                        value=0.8,
+                        step=0.05,
+                        label="Strength (0=no change, 1=full change)"
+                    )
+                    
+                    img2img_seed = gr.Number(
+                        value=-1,
+                        label="Seed (-1 = random)",
+                        precision=0
+                    )
+                    
+                    img2img_generate = gr.Button(
+                        "🖼️ Generate Image",
+                        variant="primary",
+                        size="lg"
+                    )
+                
+                with gr.Column(scale=1):
+                    img2img_output = gr.Image(label="Generated Image", type="pil")
+                    img2img_status = gr.Textbox(interactive=False, label="Result")
+            
+            img2img_generate.click(
+                generate_img2img_ui,
+                inputs=[
+                    img2img_prompt,
+                    img2img_input,
+                    img2img_negative,
+                    img2img_steps,
+                    img2img_guidance,
+                    img2img_strength,
+                    img2img_seed
+                ],
+                outputs=[img2img_output, img2img_status]
+            )
+        
+        # ===== TAB 3: INFO =====
+        with gr.Tab("ℹ️ Info"):
+            gr.Markdown(f"""
+            ## Model Information
+            - **Model**: {MODEL_CONFIG['model_id']}
+            - **Device**: CPU (HuggingFace Spaces)
+            - **Data Type**: {DTYPE}
+            - **Features**: Text-to-Image, Image-to-Image
+            
+            ## Tips
+            - ⏱️ CPU inference slower than GPU, expect 2-5 min per image
+            - 🎯 Higher steps = better quality but slower
+            - 📐 Larger dimensions = slower generation
+            - 💡 Detailed prompts = better results
+            - 🔄 Strength 0.8-0.9 recommended for img2img
+            """)
+
+# ===================== STARTUP =====================
+
+def on_load():
+    """Initialize model on startup"""
+    logger.info("🚀 Starting FLUX.1 App...")
+    pipeline_manager.initialize()
+    
+    if pipeline_manager.model_loaded:
+        return "✅ Model loaded! Ready to generate."
+    else:
+        return "⚠️ Model initialization in progress or fallback mode active."
+
+# Load model on app start
+gr.on_load(on_load)
+
+# ===================== LAUNCH =====================
+
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=True,
+        show_error=True,
+        debug=False
+    )