Changes

app.py

@@ -14,13 +14,10 @@ from huggingface_hub import login
 login(token='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')
 
 
-# model_path = WEIGHTS_DIR             # If you want to use previously trained model saved in gdrive, replace this with the full path of model in gdrive
-# headers = {"Authorization": "Bearer xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"}
-
 pipe = StableDiffusionPipeline.from_pretrained("Divyanshu04/Finetuned-model", safety_checker=None, torch_dtype=torch.float32).to("cpu")
 
 pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
-# pipe.enable_xformers_memory_efficient_attention()
+# pipe.enable_xformers_memory_efficient_attention()  #if gpu is available
 g_cuda = None
          
 FILE = Path(__file__).resolve()
@@ -53,7 +50,6 @@ def generate():
         seed = 52362
         g_cuda.manual_seed(seed)
 
-        # commandline_args = os.environ.get('COMMANDLINE_ARGS', "--skip-torch-cuda-test --no-half")
 
         with autocast("cpu"), torch.inference_mode():
             images = pipe(
@@ -67,7 +63,7 @@ def generate():
                 generator=g_cuda
             ).images
         
-        return {"message": "successful"}
+        st.image(images)
     
     else:
         st.write('<Enter parameters to generate image>')