Update main.py

main.py

@@ -1,75 +1,41 @@
+from fastapi import FastAPI, HTTPException
 from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
-from fastapi.responses import HTMLResponse
-from pydantic import BaseModel
-import gradio as gr
-import random
-import string
-import time
-from queue import Queue
-from threading import Thread
+from fastapi.responses import FileResponse
+import os
+import requests
 
 app = FastAPI()
-
-# Mount the static directory for serving HTML and CSS files
-app.mount("/static", StaticFiles(directory="static"), name="static")
-
-queue = Queue()
-queue_threshold = 100
-
-def restart_script_periodically():
-    while True:
-        random_time = random.randint(540, 600)
-        time.sleep(random_time)
-        os.execl(sys.executable, sys.executable, *sys.argv)
-
-restart_thread = Thread(target=restart_script_periodically, daemon=True)
-restart_thread.start()
-
-def add_random_noise(prompt, noise_level=0.00):
-    if noise_level == 0:
-        noise_level = 0.00
-    percentage_noise = noise_level * 5
-    num_noise_chars = int(len(prompt) * (percentage_noise / 100))
-    noise_indices = random.sample(range(len(prompt)), num_noise_chars)
-    prompt_list = list(prompt)
-    noise_chars = list(string.ascii_letters + string.punctuation + ' ' + string.digits)
-    noise_chars.extend(['😍', '💩', '😂', '🤔', '😊', '🤗', '😭', '🙄', '😷', '🤯', '🤫', '🥴', '😴', '🤩', '🥳', '😔', '😩', '🤪', '😇', '🤢', '😈', '👹', '👻', '🤖', '👽', '💀', '🎃', '🎅', '🎄', '🎁', '🎂', '🎉', '🎈', '🎊', '🎮', '❤️', '💔', '💕', '💖', '💗', '🐶', '🐱', '🐭', '🐹', '🦊', '🐻', '🐨', '🐯', '🦁', '🐘', '🔥', '🌧️', '🌞', '🌈', '💥', '🌴', '🌊', '🌺', '🌻', '🌸', '🎨', '🌅', '🌌', '☁️', '⛈️', '❄️', '☀️', '🌤️', '⛅️', '🌥️', '🌦️', '🌧️', '🌩️', '🌨️', '🌫️', '☔️', '🌬️', '💨', '🌪️', '🌈'])
-    for index in noise_indices:
-        prompt_list[index] = random.choice(noise_chars)
-    return "".join(prompt_list)
-
-# Define FastAPI models
-class PromptInput(BaseModel):
-    prompt: str
-    noise_level: float
-
-# Define FastAPI endpoints
-@app.post("/get_prompt/")
-async def get_prompt(prompt_input: PromptInput):
-    text_gen = gr.Interface.load("models/Gustavosta/MagicPrompt-Stable-Diffusion")
-    prompt_text = text_gen("dreamlikeart, " + prompt_input.prompt)
-    return {"prompt": prompt_text}
-
-@app.post("/generate_image/")
-async def generate_image(prompt_input: PromptInput):
-    proc1 = gr.Interface.load("models/dreamlike-art/dreamlike-diffusion-1.0", temp_files_path="static/tmp")
-    prompt_with_noise = add_random_noise(prompt_input.prompt, prompt_input.noise_level)
-    while queue.qsize() >= queue_threshold:
-        time.sleep(2)
-    queue.put(prompt_with_noise)
-    output_dict = proc1(prompt_with_noise)
-    output1 = output_dict['output'][0]
-    return {"image": output1}
-
-# Serve the HTML frontend
-@app.get("/", response_class=HTMLResponse)
-async def serve_frontend():
-    with open("static/index.html", "r") as file:
-        html_content = file.read()
-    return HTMLResponse(content=html_content)
-
-# Run the FastAPI server
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+API_URL = "https://ashrafb-dreamlikeart-diffusion-1-0.hf.space/"
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+
+def make_prediction(prompt, noise_level=0.0, fn_index=0):
+    headers={"Authorization": f"Bearer {HF_TOKEN}"}
+    data = {"prompt": prompt, "noise_level": noise_level}
+    response = requests.post(API_URL, headers=headers, json=data)
+    if response.status_code == 200:
+        return response.json()
+    else:
+        raise HTTPException(status_code=response.status_code, detail=response.text)
+
+@app.get("/short-prompt/")
+async def short_prompt(prompt: str):
+    try:
+        result = make_prediction(prompt)
+        return {"result": result}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/long-prompt/")
+async def long_prompt(prompt: str, noise_level: float = 0.0):
+    try:
+        result = make_prediction(prompt, noise_level, fn_index=1)
+        return {"result": result}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+app.mount("/", StaticFiles(directory="static", html=True), name="static")
+
+@app.get("/")
+def index() -> FileResponse:
+    return FileResponse(path="/app/static/index.html", media_type="text/html")