Update src/pipeline.py

src/pipeline.py

@@ -23,7 +23,7 @@ def empty_cache():
     torch.cuda.empty_cache()
     torch.cuda.reset_max_memory_allocated()
     torch.cuda.reset_peak_memory_stats()
-    print(f"Flush took: {time.time() - start}")
+    # print(f"Flush took: {time.time() - start}")
 
 def load_pipeline() -> Pipeline:    
     empty_cache()
@@ -49,7 +49,7 @@ def load_pipeline() -> Pipeline:
     #quantize_(vae, int8_weight_only())
 
     model = FluxTransformer2DModel.from_pretrained(
-        "/home/sandbox/.cache/huggingface/hub/models--RobertML--FLUX.1-schnell-int8wo/snapshots/307e0777d92df966a3c0f99f31a6ee8957a9857a", torch_dtype=dtype, use_safetensors=False
+        "/root/.cache/huggingface/hub/models--RobertML--FLUX.1-schnell-int8wo/snapshots/307e0777d92df966a3c0f99f31a6ee8957a9857a", torch_dtype=dtype, use_safetensors=False
     )
     empty_cache()
     pipeline = DiffusionPipeline.from_pretrained(
@@ -72,10 +72,6 @@ from datetime import datetime
 @torch.inference_mode()
 def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
     empty_cache()
-    try:
-        generator = Generator(pipeline.device).manual_seed(request.seed)
-        image=pipeline(request.prompt,generator=generator, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256, height=request.height, width=request.width, output_type="pil").images[0]
-    except:
-        image = img.open("./RobertML.png")
-        pass
+    generator = Generator(pipeline.device).manual_seed(request.seed)
+    image=pipeline(request.prompt,generator=generator, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256, height=request.height, width=request.width, output_type="pil").images[0]
     return(image)