AMfeta99 commited on
Commit
2baa37f
·
verified ·
1 Parent(s): 63f528a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -26
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