Update ImageGenerator.py

ImageGenerator.py

@@ -158,8 +158,30 @@ def extract_assistant_response(decoded_text):
 import torch
 from diffusers import FluxPipeline
 
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
-pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
+# pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+# pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
+# def generate_image(prompt):
+#     print('GENERATING IMAGE')
+#     prompt =extract_assistant_response(prompt)# "A cat holding a sign that says hello world"
+#     image = pipe(
+#         prompt,
+#         height=1024,
+#         width=1024,
+#         guidance_scale=3.5,
+#         num_inference_steps=50,
+#         max_sequence_length=512,
+#         generator=torch.Generator("cpu").manual_seed(0)
+#     ).images[0]
+#     #image.save("flux-dev.png")
+#     return image
+
+
+from diffusers import DiffusionPipeline
+import torch
+
+pipeline = DiffusionPipeline.from_pretrained("segmind/tiny-sd", torch_dtype=torch.float16)
+# prompt = "Portrait of a pretty girl"
+# image = pipeline(prompt).images[0]
 def generate_image(prompt):
     print('GENERATING IMAGE')
     prompt =extract_assistant_response(prompt)# "A cat holding a sign that says hello world"