Spaces:

Valtry
/

My-Image

Sleeping

App Files Files Community

Valtry commited on Aug 8, 2025

Commit

2f0bcde

verified ·

1 Parent(s): 4b1ddd1

Update app.py

Browse files

Files changed (1) hide show

app.py +132 -44

app.py CHANGED Viewed

@@ -1,64 +1,152 @@
 from fastapi import FastAPI, Request
-from fastapi.responses import HTMLResponse
 from pydantic import BaseModel
-from diffusers import StableDiffusionPipeline
 import torch
-import base64
-from io import BytesIO
-from PIL import Image
-app = FastAPI()
-# Load model
-model_id = "runwayml/stable-diffusion-v1-5"
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cpu")  # use "cuda" if GPU
-# Request body schema
-class GenerateRequest(BaseModel):
     prompt: str
-# API route for text-to-image
 @app.post("/generate")
-def generate_image(request: GenerateRequest):
-    image = pipe(request.prompt).images[0]
-    # Save to base64
-    buffer = BytesIO()
-    image.save(buffer, format="PNG")
-    img_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
-    return {"image_base64": img_str}
-# Simple HTML test page
 @app.get("/", response_class=HTMLResponse)
 async def home():
     return """
     <html>
     <head>
-        <title>Valtry Text-to-Image</title>
     </head>
-    <body style="font-family: Arial; text-align: center; padding: 50px;">
-        <h2>🖌️ Valtry Text-to-Image Generator</h2>
-        <input id="prompt" type="text" placeholder="Enter your prompt..." style="width: 300px; padding: 5px;" />
-        <button onclick="generate()">Generate</button>
-        <div id="result" style="margin-top: 20px;"></div>
-        <script>
-            async function generate() {
-                const prompt = document.getElementById("prompt").value;
-                document.getElementById("result").innerHTML = "⏳ Generating...";
-                const res = await fetch('/generate', {
-                    method: 'POST',
-                    headers: { 'Content-Type': 'application/json' },
-                    body: JSON.stringify({ prompt })
-                });
-                if (!res.ok) {
-                    document.getElementById("result").innerHTML = "❌ Error: " + res.status;
-                    return;
-                }
-                const data = await res.json();
-                document.getElementById("result").innerHTML =
-                    `<img src="data:image/png;base64,${data.image_base64}" style="max-width:100%;"/>`;
             }
-        </script>
     </body>
     </html>
-    """

+import os
+import time
+import uuid
+from io import BytesIO
+from typing import Optional
 from fastapi import FastAPI, Request
+from fastapi.responses import HTMLResponse, JSONResponse
+from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
+# Put the HF / transformers cache into a writable folder
+os.environ["HF_HOME"] = "/app/cache"
+os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
+os.makedirs("/app/cache", exist_ok=True)
+os.makedirs("/app/static", exist_ok=True)
+# Import after setting env
 import torch
+from diffusers import StableDiffusionPipeline
+# -------- CONFIG --------
+MODEL_ID = "runwayml/stable-diffusion-v1-5"   # change if you prefer another model
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+STATIC_FOLDER = "/app/static"
+# ------------------------
+app = FastAPI(title="Valtry Text→Image API")
+# Mount static folder so generated images are publicly accessible at /static/...
+app.mount("/static", StaticFiles(directory=STATIC_FOLDER), name="static")
+# Load model once at startup
+print("Loading model", MODEL_ID, "to device", DEVICE)
+# If your model is gated, set use_auth_token=os.getenv("HF_TOKEN") in from_pretrained(...)
+pipe = StableDiffusionPipeline.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
+)
+pipe = pipe.to(DEVICE)
+pipe.safety_checker = pipe.safety_checker if hasattr(pipe, "safety_checker") else None
+print("Model loaded")
+class GenerateReq(BaseModel):
     prompt: str
+    num_inference_steps: Optional[int] = 25
+    guidance_scale: Optional[float] = 7.5
+    seed: Optional[int] = None
 @app.post("/generate")
+async def generate(req: GenerateReq):
+    if not req.prompt or not req.prompt.strip():
+        return JSONResponse({"error": "prompt is required"}, status_code=400)
+    seed = req.seed if req.seed is not None else int(time.time() * 1000) % 2**32
+    generator = torch.Generator(device=DEVICE).manual_seed(seed) if DEVICE == "cuda" else None
+    # Generate image
+    try:
+        result = pipe(
+            req.prompt,
+            num_inference_steps=int(req.num_inference_steps),
+            guidance_scale=float(req.guidance_scale),
+            generator=generator,
+        )
+    except Exception as e:
+        return JSONResponse({"error": f"model generation failed: {str(e)}"}, status_code=500)
+    image = result.images[0]
+    # Save image file
+    filename = f"img_{int(time.time())}_{uuid.uuid4().hex[:8]}.png"
+    file_path = os.path.join(STATIC_FOLDER, filename)
+    image.save(file_path)
+    # Return the public URL (relative)
+    return {"url": f"/static/{filename}", "filename": filename}
+# Simple home page for quick testing in browser
 @app.get("/", response_class=HTMLResponse)
 async def home():
     return """
+    <!doctype html>
     <html>
     <head>
+      <meta charset="utf-8"/>
+      <title>Valtry Text→Image</title>
+      <style>
+        body{font-family:Arial,sans-serif;margin:32px;background:#f7f7f7}
+        input, button, textarea{font-size:16px;padding:10px;width:100%;box-sizing:border-box;margin-top:8px}
+        #result{margin-top:20px}
+        img{max-width:100%;border:1px solid #ccc;padding:6px;background:#fff}
+        label{font-weight:600}
+      </style>
     </head>
+    <body>
+      <h2>Valtry — Text → Image</h2>
+      <label>Prompt</label>
+      <textarea id="prompt" rows="3" placeholder="A fantasy castle on a cliff at sunset"></textarea>
+      <label>Max steps (num_inference_steps)</label>
+      <input id="steps" type="number" value="25" min="1" max="150"/>
+      <label>Guidance scale</label>
+      <input id="scale" type="number" value="7.5" step="0.1" min="1" max="20"/>
+      <label>Seed (optional)</label>
+      <input id="seed" type="number" placeholder="leave empty for random"/>
+      <button onclick="generate()">Generate Image</button>
+      <div id="status"></div>
+      <div id="result"></div>
+      <script>
+        async function generate(){
+          const prompt = document.getElementById('prompt').value;
+          const steps = parseInt(document.getElementById('steps').value || 25);
+          const scale = parseFloat(document.getElementById('scale').value || 7.5);
+          const seedVal = document.getElementById('seed').value;
+          document.getElementById('status').textContent = "⏳ Generating — this may take a bit...";
+          document.getElementById('result').innerHTML = "";
+          const body = {
+            prompt: prompt,
+            num_inference_steps: steps,
+            guidance_scale: scale
+          };
+          if (seedVal) body.seed = parseInt(seedVal);
+          try {
+            const res = await fetch('/generate', {
+              method: 'POST',
+              headers: {'Content-Type': 'application/json'},
+              body: JSON.stringify(body)
+            });
+            if (!res.ok) {
+              const txt = await res.text();
+              document.getElementById('status').textContent = '❌ Error ' + res.status + ': ' + txt;
+              return;
             }
+            const data = await res.json();
+            document.getElementById('status').textContent = '✅ Done — image below';
+            document.getElementById('result').innerHTML = `<img src="${data.url}" alt="generated-image"/>`;
+          } catch (err) {
+            document.getElementById('status').textContent = '❌ Exception: ' + err.message;
+          }
+        }
+      </script>
     </body>
     </html>
+    """
+# optional health route
+@app.get("/health")
+async def health():
+    return {"status": "ok"}