Upload 3 files

Dockerfile

@@ -0,0 +1,18 @@
+# Use an official Python runtime as a parent image
+FROM python:3.9-slim
+
+# Set the working directory in the container
+WORKDIR /app
+
+# Copy the requirements file and install dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the FastAPI application code into the container
+COPY main.py .
+
+# Expose port 80 for the app to be accessible externally
+EXPOSE 7860
+
+# Run the application with Uvicorn
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

main.py

@@ -0,0 +1,115 @@
+import os
+import json
+import requests
+from fastapi import FastAPI, HTTPException, Form
+from fastapi.responses import FileResponse
+
+app = FastAPI()
+
+def generate_image_and_download(prompt: str, id_image: str) -> str:
+    # Base API endpoint for image generation
+    base_url = "https://yanze-pulid-flux.hf.space/call/generate_image"
+    
+    # JSON payload with 13 custom parameters
+    payload = {
+        "data": [
+            prompt,   # Prompt text
+            {"path": id_image},  # ID image file
+            0,        # start_step
+            4,        # guidance
+            "-1",     # seed ("-1" for random)
+            1,        # true_cfg
+            896,      # width
+            1152,     # height
+            20,       # num_steps
+            1,        # id_weight
+            "bad quality, worst quality, text, signature, watermark, extra limbs",  # neg_prompt
+            1,        # timestep_to_start_cfg
+            128       # max_sequence_length
+        ]
+    }
+    headers = {"Content-Type": "application/json"}
+    
+    # Send the initial POST request
+    response = requests.post(base_url, json=payload, headers=headers)
+    if response.status_code != 200:
+        raise HTTPException(status_code=response.status_code, detail=f"Error in POST request: {response.text}")
+    
+    # Extract event ID from the response (using splitting as in your shell command)
+    try:
+        event_id = response.text.split('"')[3]
+    except IndexError:
+        raise HTTPException(status_code=500, detail="Failed to extract event ID from response.")
+    
+    # Poll the event endpoint to get the streaming response
+    poll_url = f"{base_url}/{event_id}"
+    complete_event_received = False
+    data_buffer = []
+    
+    try:
+        with requests.get(poll_url, stream=True) as poll_response:
+            if poll_response.status_code != 200:
+                raise HTTPException(status_code=poll_response.status_code, detail=f"Error while polling event: {poll_response.text}")
+            for line in poll_response.iter_lines(decode_unicode=True):
+                if not line:
+                    if complete_event_received and data_buffer:
+                        break
+                    data_buffer = []
+                    continue
+
+                if line.startswith("event:"):
+                    event_type = line[len("event:"):].strip()
+                    if event_type == "complete":
+                        complete_event_received = True
+                elif line.startswith("data:"):
+                    data_line = line[len("data:"):].strip()
+                    if complete_event_received:
+                        data_buffer.append(data_line)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error while polling: {e}")
+
+    if not data_buffer:
+        raise HTTPException(status_code=500, detail="No complete event data received.")
+
+    complete_event_data = "\n".join(data_buffer)
+    try:
+        json_data = json.loads(complete_event_data)
+    except json.JSONDecodeError as e:
+        raise HTTPException(status_code=500, detail=f"JSON decode error: {e}")
+
+    # Extract the image path from the JSON response.
+    if isinstance(json_data, list) and len(json_data) > 0 and isinstance(json_data[0], dict) and "path" in json_data[0]:
+        image_path = json_data[0]["path"]
+    else:
+        raise HTTPException(status_code=500, detail="Could not extract image path from response JSON.")
+    
+    # Construct the download URL
+    download_base_url = "https://yanze-pulid-flux.hf.space/file="
+    download_url = f"{download_base_url}{image_path}"
+    
+    # Download the image file
+    try:
+        image_response = requests.get(download_url, stream=True)
+        if image_response.status_code == 200:
+            # Use the original filename from the image path
+            filename = image_path.split("/")[-1]
+            with open(filename, "wb") as f:
+                for chunk in image_response.iter_content(chunk_size=1024):
+                    f.write(chunk)
+            return filename
+        else:
+            raise HTTPException(status_code=image_response.status_code, detail="Failed to download image.")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error downloading image: {e}")
+
+@app.post("/generate")
+async def generate(prompt: str = Form(...), id_image: str = Form(...)):
+    """
+    Expects form data with 'prompt' and 'id_image' keys.
+    When complete, returns the generated image file.
+    """
+    filename = generate_image_and_download(prompt, id_image)
+    if os.path.exists(filename):
+        return FileResponse(path=filename, filename=filename, media_type="application/octet-stream")
+    else:
+        raise HTTPException(status_code=500, detail="Image file not found after generation.")

requirements.txt

@@ -0,0 +1,3 @@
+fastapi
+uvicorn
+requests