Normal

app.py

@@ -1,108 +1,55 @@
-# import gradio as gr
-# import torch
-# from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
-# from PIL import Image
-# import base64
-# from io import BytesIO
-
-# # You can change these:
-# BASE_MODEL = "runwayml/stable-diffusion-v1-5"
-# CONTROLNET_ID = "lllyasviel/sd-controlnet-canny"  # placeholder; change to a QR-focused ControlNet if you have one
-
-# device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# controlnet = ControlNetModel.from_pretrained(
-#     CONTROLNET_ID, torch_dtype=torch.float16 if device=="cuda" else torch.float32
-# )
-
-# pipe = StableDiffusionControlNetPipeline.from_pretrained(
-#     BASE_MODEL,
-#     controlnet=controlnet,
-#     torch_dtype=torch.float16 if device=="cuda" else torch.float32,
-#     safety_checker=None
-# )
-# pipe.to(device)
-
-# def generate(prompt, control_image, guidance_scale=7.5, steps=30, seed=0):
-#     print("API called:", type(control_image))
-#     generator = torch.Generator(device=device).manual_seed(int(seed)) if seed else None
-#     img = pipe(
-#         prompt=prompt,
-#         image=control_image,
-#         num_inference_steps=int(steps),
-#         guidance_scale=float(guidance_scale),
-#         generator=generator
-#     ).images[0]
-#     return img
-
-# with gr.Blocks() as demo:
-#     gr.Markdown("# ControlNet Image Generator")
-#     with gr.Row():
-#         prompt = gr.Textbox(label="Prompt", value="A futuristic poster, high detail")
-#         seed = gr.Number(label="Seed (0=random)", value=0)
-#     with gr.Row():
-#         control = gr.Image(type="pil", label="Control image (e.g., QR or edge map)")
-#         steps = gr.Slider(10, 50, 30, step=1, label="Steps")
-#         guidance = gr.Slider(1.0, 12.0, 7.5, step=0.1, label="Guidance scale")
-#     out = gr.Image(label="Result")
-
-#     btn = gr.Button("Generate")
-#     btn.click(generate, [prompt, control, guidance, steps, seed], out)
-
-#     # Enable simple API use
-#     gr.Examples([], inputs=[prompt, control, guidance, steps, seed], outputs=out)
-
-# demo.launch()
-
-# app.py  — minimal, API-stable echo server + Gradio UI
-import base64, io
-from typing import List, Any
-from PIL import Image
 import gradio as gr
-from fastapi import FastAPI
-from pydantic import BaseModel
-
-# ---- helpers ----
-def data_url_to_pil(data_url: str) -> Image.Image:
-    # expects "data:image/<ext>;base64,<...>"
-    b64 = data_url.split(",", 1)[1]
-    return Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
+import torch
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from PIL import Image
+import base64
+from io import BytesIO
 
-def pil_to_data_url(img: Image.Image) -> str:
-    buf = io.BytesIO()
-    img.save(buf, format="PNG")
-    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
+# You can change these:
+BASE_MODEL = "runwayml/stable-diffusion-v1-5"
+CONTROLNET_ID = "lllyasviel/sd-controlnet-canny"  # placeholder; change to a QR-focused ControlNet if you have one
 
-# ---- your model fn (for now just echo control image) ----
-def generate(prompt: str, control_image: Image.Image, guidance: float, steps: int, seed: int):
-    # TODO: replace with ControlNet pipeline call
-    return control_image
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# ---- Gradio UI (not required for API, but nice to see) ----
-demo = gr.Interface(
-    fn=generate,
-    inputs=[
-        gr.Textbox(label="Prompt", value="High-quality artistic style, preserve structure"),
-        gr.Image(type="pil", label="Control Image (PNG/JPG)"),
-        gr.Slider(1, 12, 7.5, step=0.1, label="Guidance scale"),
-        gr.Slider(10, 50, 30, step=1, label="Steps"),
-        gr.Number(0, label="Seed"),
-    ],
-    outputs=gr.Image(label="Result"),
+controlnet = ControlNetModel.from_pretrained(
+    CONTROLNET_ID, torch_dtype=torch.float16 if device=="cuda" else torch.float32
 )
 
-# ---- FastAPI endpoint (stable path for Postman/Next.js) ----
-app = FastAPI()
-
-class PredictIn(BaseModel):
-    data: List[Any]  # [prompt, controlImageBase64, guidance, steps, seed]
-
-@app.post("/api/predict/0")
-def predict(payload: PredictIn):
-    prompt, control_b64, guidance, steps, seed = payload.data
-    img = data_url_to_pil(control_b64) if isinstance(control_b64, str) else control_b64
-    out = generate(str(prompt), img, float(guidance), int(steps), int(seed))
-    return { "data": [ pil_to_data_url(out) ] }
-
-# mount Gradio at "/"
-app = gr.mount_gradio_app(app, demo, path="/")
+pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    BASE_MODEL,
+    controlnet=controlnet,
+    torch_dtype=torch.float16 if device=="cuda" else torch.float32,
+    safety_checker=None
+)
+pipe.to(device)
+
+def generate(prompt, control_image, guidance_scale=7.5, steps=30, seed=0):
+    print("API called:", type(control_image))
+    generator = torch.Generator(device=device).manual_seed(int(seed)) if seed else None
+    img = pipe(
+        prompt=prompt,
+        image=control_image,
+        num_inference_steps=int(steps),
+        guidance_scale=float(guidance_scale),
+        generator=generator
+    ).images[0]
+    return img
+
+with gr.Blocks() as demo:
+    gr.Markdown("# ControlNet Image Generator")
+    with gr.Row():
+        prompt = gr.Textbox(label="Prompt", value="A futuristic poster, high detail")
+        seed = gr.Number(label="Seed (0=random)", value=0)
+    with gr.Row():
+        control = gr.Image(type="pil", label="Control image (e.g., QR or edge map)")
+        steps = gr.Slider(10, 50, 30, step=1, label="Steps")
+        guidance = gr.Slider(1.0, 12.0, 7.5, step=0.1, label="Guidance scale")
+    out = gr.Image(label="Result")
+
+    btn = gr.Button("Generate")
+    btn.click(generate, [prompt, control, guidance, steps, seed], out)
+
+    # Enable simple API use
+    gr.Examples([], inputs=[prompt, control, guidance, steps, seed], outputs=out)
+
+demo.launch()