Update app.py
Browse files
app.py
CHANGED
@@ -91,7 +91,7 @@ class TextToImageTool(Tool):
|
|
91 |
def forward(self, prompt):
|
92 |
return self.client.text_to_image(prompt)
|
93 |
'''
|
94 |
-
|
95 |
class TextToImageTool(Tool):
|
96 |
description = "This tool creates an image according to a prompt. Add details like 'high-res, photorealistic'."
|
97 |
name = "image_generator"
|
@@ -116,32 +116,7 @@ class TextToImageTool(Tool):
|
|
116 |
def forward(self, prompt):
|
117 |
image = self.pipe(prompt).images[0]
|
118 |
return image
|
119 |
-
'''
|
120 |
-
|
121 |
-
# Use a public model (you can swap for a gated model with token later)
|
122 |
-
pipe = DiffusionPipeline.from_pretrained(
|
123 |
-
"stabilityai/stable-diffusion-2-1",
|
124 |
-
torch_dtype=torch.float16
|
125 |
-
).to("cuda" if torch.cuda.is_available() else "cpu")
|
126 |
-
|
127 |
-
# --------------------------------------------------------
|
128 |
-
# Define a simple image generation function
|
129 |
-
# --------------------------------------------------------
|
130 |
-
|
131 |
-
def generate_image(prompt: str) -> Image.Image:
|
132 |
-
"""Generate an image from a text prompt using a diffusion model."""
|
133 |
-
image = pipe(prompt).images[0]
|
134 |
-
return image
|
135 |
-
|
136 |
-
# --------------------------------------------------------
|
137 |
-
# Wrap it into a Tool using from_function
|
138 |
-
# --------------------------------------------------------
|
139 |
|
140 |
-
image_generation_tool = Tool.from_function(
|
141 |
-
name="image_generator",
|
142 |
-
description="Generate an image from a prompt, e.g. 'a dog on Mars in high-res'.",
|
143 |
-
func=generate_image
|
144 |
-
)
|
145 |
|
146 |
# =========================================================
|
147 |
# Tool and Agent Initialization
|
|
|
91 |
def forward(self, prompt):
|
92 |
return self.client.text_to_image(prompt)
|
93 |
'''
|
94 |
+
|
95 |
class TextToImageTool(Tool):
|
96 |
description = "This tool creates an image according to a prompt. Add details like 'high-res, photorealistic'."
|
97 |
name = "image_generator"
|
|
|
116 |
def forward(self, prompt):
|
117 |
image = self.pipe(prompt).images[0]
|
118 |
return image
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
119 |
|
|
|
|
|
|
|
|
|
|
|
120 |
|
121 |
# =========================================================
|
122 |
# Tool and Agent Initialization
|