Checkconnect

app.py

@@ -1,55 +1,108 @@
-import gradio as gr
-import torch
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+# 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 base64
-from io import BytesIO
+import gradio as gr
+from fastapi import FastAPI
+from pydantic import BaseModel
 
-# 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
+# ---- 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")
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+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()
 
-controlnet = ControlNetModel.from_pretrained(
-    CONTROLNET_ID, torch_dtype=torch.float16 if device=="cuda" else torch.float32
-)
+# ---- 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
 
-pipe = StableDiffusionControlNetPipeline.from_pretrained(
-    BASE_MODEL,
-    controlnet=controlnet,
-    torch_dtype=torch.float16 if device=="cuda" else torch.float32,
-    safety_checker=None
+# ---- 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"),
 )
-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()
+
+# ---- 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="/")