Update main.py

main.py

@@ -2,10 +2,7 @@ from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import FileResponse
 
-import platform
 import subprocess
-import logging
-import urllib.request
 import os
 import json
 import uuid
@@ -18,23 +15,16 @@ app = FastAPI()
 @app.get("/generate")
 def generate_image(prompt, inference_steps, model):
     torch.cuda.empty_cache()
-    print(f"Is CUDA available: {torch.cuda.is_available()}")
 
     pipeline = StableDiffusionPipeline.from_pretrained(str(model), torch_dtype=torch.float16)
-    
-    #pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
-    #another comment
     pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
-    
     pipeline = pipeline.to("cuda")
+ 
     image = pipeline(prompt, num_inference_steps=int(inference_steps), height=512, width=512).images[0]
 
     filename = str(uuid.uuid4()) + ".jpg"
     image.save(filename)
     
-    print(filename)
-    
-    # Data to be written
     assertion = {
                "assertions": [
                    {
@@ -48,44 +38,16 @@ def generate_image(prompt, inference_steps, model):
                ]
                }
 
-    print(assertion)
-    
     json_object = json.dumps(assertion, indent=4)
     with open("assertion.json", "w") as outfile:
        outfile.write(json_object)
 
     subprocess.check_output(['./truepic-sign', 'init', 'file-system', '--api-key', os.environ.get("api_key")])
     subprocess.check_output(['./truepic-sign', 'sign', filename, '--assertions-file', 'assertion.json', '--output', (os.getcwd() + '/static/' + filename)])
-    # subprocess.check_output(['./truepic-sign', 'sign', filename, '--output', (os.getcwd() + '/static/' + filename)])
-
+    
     return {"response": filename}
 
 
-@app.get("/generate-picsum")
-def generate_picsum(prompt):
-    local_filename, headers = urllib.request.urlretrieve(('https://picsum.photos/id/' + prompt + '/800/800'))
-
-    # Data to be written
-    assertion = {
-                "assertions": [
-                    {
-                    "label": "com.truepic.custom.ai",
-                    "data": {
-                        "model_name": "Picsum",
-                        "model_version": "1.0",
-                        "prompt": prompt
-                    }
-                    }
-                ]
-                }
-    
-    json_object = json.dumps(assertion, indent=4)
-    with open("assertion.json", "w") as outfile:
-        outfile.write(json_object)
-    
-    subprocess.check_output(['./truepic-sign', 'sign', local_filename, '--profile', 'demo', '--assertions', 'assertion.json', '--output', (os.getcwd() + '/static/output.jpg')])
-    return {"response": "success"}
-
 app.mount("/", StaticFiles(directory="static", html=True), name="static")
 
 @app.get("/")