Update main.py

main.py

@@ -1,68 +1,56 @@
-from fastapi import FastAPI, HTTPException, Request
-from rdkit import Chem
-from rdkit.Chem import Descriptors
-from rdkit.Chem import rdMolDescriptors
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
 import requests
-from fastapi.middleware.cors import CORSMiddleware
+from urllib.parse import quote
+import asyncio
+from io import BytesIO
 
 app = FastAPI()
 
-app.add_middleware(
-    CORSMiddleware,
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-def name_to_smiles(name: str) -> str:
-    url = f"https://cactus.nci.nih.gov/chemical/structure/{name}/smiles"
-    response = requests.get(url)
-    if response.status_code == 200:
-        return response.text
-    return None
-
-def get_molecule_info(mol):
-    mol_weight = Descriptors.MolWt(mol)
-    num_atoms = mol.GetNumAtoms()
-    num_bonds = mol.GetNumBonds()
-    mol_formula = rdMolDescriptors.CalcMolFormula(mol)
-    tpsa = Descriptors.TPSA(mol)
-    mol_logp = Descriptors.MolLogP(mol)
-    num_rotatable_bonds = Descriptors.NumRotatableBonds(mol)
-
-    return {
-        'molecular_weight': mol_weight,
-        'number_of_atoms': num_atoms,
-        'number_of_bonds': num_bonds,
-        'molecular_formula': mol_formula,
-        'tpsa': tpsa,
-        'logP': mol_logp,
-        'number_of_rotatable_bonds': num_rotatable_bonds
+def create_job(prompt, model, sampler, seed, neg):
+    if model is None:
+        model = 'Realistic_Vision_V5.0.safetensors [614d1063]'
+    if sampler is None:
+        sampler = 'DPM++ 2M Karras'
+    if seed is None:
+        seed = '-1'
+    if neg is None:
+        neg = "(long list of negative prompts removed for brevity)"
+    url = 'https://api.prodia.com/generate'
+    params = {
+        'new': 'true',
+        'prompt': quote(prompt),
+        'model': model,
+        'negative_prompt': neg,
+        'steps': '100',
+        'cfg': '9.5',
+        'seed': seed,
+        'sampler': sampler,
+        'upscale': 'True',
+        'aspect_ratio': 'square'
     }
-
-@app.get("/")
-def home():
-    return "Hello! This is ChemDB."
-
-
-@app.get("/molecule_info/{name}")
-async def read_molecule_info(name: str, request: Request):
-    if not name:
-        raise HTTPException(status_code=400, detail="No molecule name provided")
-
-    smiles = name_to_smiles(name)
-    if not smiles:
-        raise HTTPException(status_code=400, detail="Could not fetch SMILES string for provided name")
-
-    mol = Chem.MolFromSmiles(smiles)
-
-    if not mol:
-        raise HTTPException(status_code=400, detail="Molecule not recognized")
-
-    info = get_molecule_info(mol)
-
-    return info
+    response = requests.get(url, params=params)
+    response.raise_for_status()
+    return response.json()['job']
+
+@app.get("/generate_image")
+async def generate_image(prompt: str, model: str = None, sampler: str = None, seed: str = None, neg: str = None):
+    job_id = create_job(prompt, model, sampler, seed, neg)
+    url = f'https://api.prodia.com/job/{job_id}'
+    headers = {'accept': '*/*'}
+    
+    while True:
+        response = requests.get(url=url, headers=headers)
+        response.raise_for_status()
+        job_response = response.json()
+        if job_response['status'] == 'succeeded':
+            image_url = f'https://images.prodia.xyz/{job_id}.png'
+            image_response = requests.get(image_url)
+            image_response.raise_for_status()
+            image = BytesIO(image_response.content)
+            return StreamingResponse(image, media_type='image/png')
+        await asyncio.sleep(2)  # Add a delay to prevent excessive requests
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=3000)
+    uvicorn.run(app, host="0.0.0.0", port=8000, debug=True)