Create app.py

app.py

@@ -0,0 +1,197 @@
+import gradio as gr
+import torch
+from diffusers import StableDiffusionPipeline
+from PIL import Image
+import io
+import os
+from datetime import datetime
+import time
+import json
+import uuid
+import random
+from fastapi import FastAPI, BackgroundTasks
+from pydantic import BaseModel
+import requests
+from huggingface_hub import HfApi
+
+# =============================================
+# CONFIGURATION
+# =============================================
+HF_TOKEN = os.environ.get("HF_TOKEN")
+HF_USERNAME = "yukee1992"
+DATASET_NAME = "video-project-images"
+DATASET_ID = f"{HF_USERNAME}/{DATASET_NAME}"
+
+print("=" * 60)
+print("🚀 STARTING IMAGE GENERATOR")
+print("=" * 60)
+print(f"📦 HF Dataset: {DATASET_ID}")
+print(f"🔑 HF Token: {'✅ Set' if HF_TOKEN else '❌ Missing'}")
+
+# Create backup directory
+BACKUP_DIR = "generated_images_backup"
+os.makedirs(BACKUP_DIR, exist_ok=True)
+
+# Initialize FastAPI
+app = FastAPI()
+
+# Global model cache
+model = None
+model_lock = threading.Lock()
+
+# =============================================
+# MODEL LOADING
+# =============================================
+def load_model():
+    global model
+    if model is None:
+        with model_lock:
+            if model is None:
+                print("🔄 Loading model...")
+                model = StableDiffusionPipeline.from_pretrained(
+                    "runwayml/stable-diffusion-v1-5",
+                    torch_dtype=torch.float32,
+                    safety_checker=None
+                ).to("cpu")
+                print("✅ Model loaded!")
+    return model
+
+# Preload model at startup
+load_model()
+
+# =============================================
+# HF DATASET FUNCTIONS
+# =============================================
+def upload_to_hf_dataset(image, project_id, scene_num):
+    """Upload image to HF Dataset"""
+    if not HF_TOKEN:
+        return None
+    
+    try:
+        # Convert image to bytes
+        img_bytes = io.BytesIO()
+        image.save(img_bytes, format='PNG')
+        img_data = img_bytes.getvalue()
+        
+        # Create filename
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        filename = f"scene_{scene_num:03d}_{timestamp}.png"
+        path_in_repo = f"data/projects/{project_id}/{filename}"
+        
+        # Upload
+        api = HfApi(token=HF_TOKEN)
+        api.upload_file(
+            path_or_fileobj=img_data,
+            path_in_repo=path_in_repo,
+            repo_id=DATASET_ID,
+            repo_type="dataset"
+        )
+        
+        url = f"https://huggingface.co/datasets/{DATASET_ID}/resolve/main/{path_in_repo}"
+        return url
+    except Exception as e:
+        print(f"❌ Upload failed: {e}")
+        return None
+
+# =============================================
+# IMAGE GENERATION
+# =============================================
+def generate_image(prompt, project_id=None, scene_num=1):
+    """Generate a single image"""
+    try:
+        pipe = load_model()
+        
+        # Generate
+        image = pipe(
+            prompt,
+            num_inference_steps=25,
+            guidance_scale=7.5,
+            generator=torch.Generator(device="cpu").manual_seed(random.randint(1, 999999))
+        ).images[0]
+        
+        # Save locally
+        local_path = os.path.join(BACKUP_DIR, f"{uuid.uuid4()}.png")
+        image.save(local_path)
+        
+        # Upload to HF Dataset if project_id provided
+        hf_url = None
+        if project_id:
+            hf_url = upload_to_hf_dataset(image, project_id, scene_num)
+        
+        return {
+            "image": image,
+            "local_path": local_path,
+            "hf_url": hf_url
+        }
+    except Exception as e:
+        print(f"❌ Generation failed: {e}")
+        raise
+
+# =============================================
+# API ENDPOINTS
+# =============================================
+class GenerateRequest(BaseModel):
+    prompt: str
+    project_id: Optional[str] = None
+
+@app.post("/generate")
+async def generate(request: GenerateRequest):
+    """Generate a single image"""
+    try:
+        result = generate_image(request.prompt, request.project_id)
+        return {
+            "status": "success",
+            "hf_url": result["hf_url"],
+            "local_path": result["local_path"]
+        }
+    except Exception as e:
+        return {"status": "error", "message": str(e)}
+
+@app.get("/health")
+async def health():
+    """Health check"""
+    return {
+        "status": "healthy",
+        "model_loaded": model is not None,
+        "hf_dataset": DATASET_ID if HF_TOKEN else "disabled"
+    }
+
+# =============================================
+# GRADIO INTERFACE (Optional UI)
+# =============================================
+if 'gradio' in globals():
+    def gradio_generate(prompt):
+        if not prompt:
+            return None
+        result = generate_image(prompt)
+        return result["image"]
+    
+    iface = gr.Interface(
+        fn=gradio_generate,
+        inputs=gr.Textbox(label="Prompt", placeholder="Enter your prompt..."),
+        outputs=gr.Image(label="Generated Image"),
+        title="Image Generator",
+        description="Generate images with Stable Diffusion"
+    )
+    
+    # Mount Gradio
+    gr.mount_gradio_app(app, iface, path="/")
+else:
+    @app.get("/")
+    async def root():
+        return {"message": "API is running. Use /generate endpoint"}
+
+# =============================================
+# MAIN
+# =============================================
+if __name__ == "__main__":
+    import uvicorn
+    print("\n" + "=" * 60)
+    print("🌐 Server starting on port 7860")
+    print(f"📊 API endpoints:")
+    print(f"   - POST /generate")
+    print(f"   - GET  /health")
+    print(f"   - GET  / (UI if gradio enabled)")
+    print("=" * 60)
+    
+    uvicorn.run(app, host="0.0.0.0", port=7860)